diff --git a/.codespellrc b/.codespellrc
index 261197eb..07d4908c 100644
--- a/.codespellrc
+++ b/.codespellrc
@@ -3,7 +3,7 @@
ignore-words-list = Skyflow,skyflow,skyflowapi,skyflowapis,deidentify,reidentify,detokenize,upsert,upserting,binlookup,byot,creds,fpe,devsecops,formdata,vaultid,dotenv,usecwd,runid,dateutil,Homogenous
# Skip these files and folders
-skip = .git,.venv,venv,env,__pycache__,*.pyc,*.egg-info,dist,build,.idea,.vscode,*.log,requirements.txt,./skyflow/generated,setup.py
+skip = .git,.venv,venv,env,__pycache__,*.pyc,*.egg-info,dist,build,.idea,.vscode,*.log,requirements.txt,generated,setup.py
# If you want to verify it is working, you can uncomment this line to see what files it checks
# count =
diff --git a/.github/workflows/beta-release.yml b/.github/workflows/beta-release.yml
index 7ad03858..61f9fa17 100644
--- a/.github/workflows/beta-release.yml
+++ b/.github/workflows/beta-release.yml
@@ -4,15 +4,28 @@ on:
push:
tags: '*.*.*b*'
paths-ignore:
- - "setup.py"
+ - "*/setup.py"
- "*.yml"
- "*.md"
- - "skyflow/utils/_version.py"
+ - "*/skyflow/utils/_version.py"
jobs:
build-and-deploy:
+ strategy:
+ matrix:
+ include:
+ - variant: v2
+ package-name: skyflow
+ tag-prefix: ''
+ - variant: flowvault
+ package-name: skyflow_flowvault
+ tag-prefix: 'flowvault-'
+ if: (matrix.variant == 'flowvault' && startsWith(github.ref_name, 'flowvault-')) || (matrix.variant == 'v2' && !startsWith(github.ref_name, 'flowvault-'))
uses: ./.github/workflows/shared-build-and-deploy.yml
with:
ref: ${{ github.ref_name }}
tag: 'beta'
+ variant: ${{ matrix.variant }}
+ package-name: ${{ matrix.package-name }}
+ tag-prefix: ${{ matrix.tag-prefix }}
secrets: inherit
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 952ccb38..e82cabd0 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -18,7 +18,39 @@ jobs:
error: 'One of your your commit messages is not matching the format with JIRA ID Ex: ( SDK-123 commit message )'
test:
+ strategy:
+ fail-fast: false
+ matrix:
+ include:
+ - variant: v2
+ package-name: skyflow
+ coverage-omit: "skyflow/generated/*,skyflow/utils/validations/*,skyflow/vault/data/*,skyflow/vault/detect/*,skyflow/vault/tokens/*,skyflow/vault/connection/*,skyflow/error/*,skyflow/utils/enums/*,skyflow/vault/controller/_audit.py,skyflow/vault/controller/_bin_look_up.py"
+ - variant: flowvault
+ package-name: skyflow_flowvault
+ coverage-omit: "skyflow_flowvault/generated/*"
uses: ./.github/workflows/shared-tests.yml
with:
python-version: '3.9'
+ variant: ${{ matrix.variant }}
+ package-name: ${{ matrix.package-name }}
+ coverage-omit: ${{ matrix.coverage-omit }}
secrets: inherit
+
+ test-common:
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v2
+ - uses: actions/setup-python@v2
+ with:
+ python-version: '3.9'
+ - run: pip install -e ./common
+ - run: pip install coverage
+ - run: python -m coverage run --source=common --omit="common/generated/*,common/tests/*" -m unittest discover -s common/tests -t .
+ - run: coverage xml -o test-coverage.xml
+ - name: Codecov
+ uses: codecov/codecov-action@v2.1.0
+ with:
+ token: ${{ secrets.CODECOV_REPO_UPLOAD_TOKEN }}
+ files: test-coverage.xml
+ name: codecov-skyflow-python-common
+ verbose: true
diff --git a/.github/workflows/internal-release.yml b/.github/workflows/internal-release.yml
index 2e273096..b91e5b41 100644
--- a/.github/workflows/internal-release.yml
+++ b/.github/workflows/internal-release.yml
@@ -5,19 +5,33 @@ on:
tags-ignore:
- '*.*'
paths-ignore:
- - "setup.py"
+ - "*/setup.py"
- "*.yml"
- "*.md"
- - "skyflow/utils/_version.py"
+ - "*/skyflow/utils/_version.py"
- "samples/**"
+ - "flowvault/samples/**"
branches:
- release/*
+ - flowvault-release/*
jobs:
build-and-deploy:
+ strategy:
+ matrix:
+ include:
+ - variant: v2
+ package-name: skyflow
+ tag-prefix: ''
+ - variant: flowvault
+ package-name: skyflow_flowvault
+ tag-prefix: 'flowvault-'
+ if: (matrix.variant == 'flowvault' && startsWith(github.ref_name, 'flowvault-')) || (matrix.variant == 'v2' && !startsWith(github.ref_name, 'flowvault-'))
uses: ./.github/workflows/shared-build-and-deploy.yml
with:
ref: ${{ github.ref_name }}
tag: 'internal'
+ variant: ${{ matrix.variant }}
+ package-name: ${{ matrix.package-name }}
+ tag-prefix: ${{ matrix.tag-prefix }}
secrets: inherit
-
\ No newline at end of file
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 01b8c040..de816ac0 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -7,7 +7,39 @@ on:
jobs:
test:
+ strategy:
+ fail-fast: false
+ matrix:
+ include:
+ - variant: v2
+ package-name: skyflow
+ coverage-omit: "skyflow/generated/*,skyflow/utils/validations/*,skyflow/vault/data/*,skyflow/vault/detect/*,skyflow/vault/tokens/*,skyflow/vault/connection/*,skyflow/error/*,skyflow/utils/enums/*,skyflow/vault/controller/_audit.py,skyflow/vault/controller/_bin_look_up.py"
+ - variant: flowvault
+ package-name: skyflow_flowvault
+ coverage-omit: "skyflow_flowvault/generated/*"
uses: ./.github/workflows/shared-tests.yml
with:
python-version: '3.9'
+ variant: ${{ matrix.variant }}
+ package-name: ${{ matrix.package-name }}
+ coverage-omit: ${{ matrix.coverage-omit }}
secrets: inherit
+
+ test-common:
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v2
+ - uses: actions/setup-python@v2
+ with:
+ python-version: '3.9'
+ - run: pip install -e ./common
+ - run: pip install coverage
+ - run: python -m coverage run --source=common --omit="common/generated/*,common/tests/*" -m unittest discover -s common/tests -t .
+ - run: coverage xml -o test-coverage.xml
+ - name: Codecov
+ uses: codecov/codecov-action@v2.1.0
+ with:
+ token: ${{ secrets.CODECOV_REPO_UPLOAD_TOKEN }}
+ files: test-coverage.xml
+ name: codecov-skyflow-python-common
+ verbose: true
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index d062daf4..01152204 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -4,15 +4,28 @@ on:
push:
tags: "*.*.*"
paths-ignore:
- - "setup.py"
+ - "*/setup.py"
- "*.yml"
- "*.md"
- - "skyflow/utils/_version.py"
+ - "*/skyflow/utils/_version.py"
jobs:
build-and-deploy:
+ strategy:
+ matrix:
+ include:
+ - variant: v2
+ package-name: skyflow
+ tag-prefix: ''
+ - variant: flowvault
+ package-name: skyflow_flowvault
+ tag-prefix: 'flowvault-'
+ if: (matrix.variant == 'flowvault' && startsWith(github.ref_name, 'flowvault-')) || (matrix.variant == 'v2' && !startsWith(github.ref_name, 'flowvault-'))
uses: ./.github/workflows/shared-build-and-deploy.yml
with:
ref: main
tag: 'public'
+ variant: ${{ matrix.variant }}
+ package-name: ${{ matrix.package-name }}
+ tag-prefix: ${{ matrix.tag-prefix }}
secrets: inherit
diff --git a/.github/workflows/semgrep.yml b/.github/workflows/semgrep.yml
index bce5fc8e..c286b921 100644
--- a/.github/workflows/semgrep.yml
+++ b/.github/workflows/semgrep.yml
@@ -20,7 +20,7 @@ jobs:
- name: Run Semgrep
run: |
- semgrep --config .semgreprules/customRule.yml --config auto --severity ERROR --sarif . > results.sarif
+ semgrep --config .semgreprules/customRule.yml --config auto --severity ERROR --exclude "**/generated/**" --sarif . > results.sarif
- name: Upload SARIF file
uses: github/codeql-action/upload-sarif@v3
diff --git a/.github/workflows/shared-build-and-deploy.yml b/.github/workflows/shared-build-and-deploy.yml
index 135b87bc..adffc3fe 100644
--- a/.github/workflows/shared-build-and-deploy.yml
+++ b/.github/workflows/shared-build-and-deploy.yml
@@ -7,12 +7,29 @@ on:
description: 'Git reference to use (e.g., main or branch name)'
required: true
type: string
-
+
tag:
description: 'Release Tag'
required: true
type: string
+ variant:
+ description: 'Build variant directory to release (v2 or flowvault)'
+ required: true
+ type: string
+
+ package-name:
+ description: 'Importable package name for this variant (e.g. skyflow or skyflow_flowvault)'
+ required: false
+ type: string
+ default: 'skyflow'
+
+ tag-prefix:
+ description: 'Prefix distinguishing this variant''s git tags from other variants'' (e.g. "flowvault-" for flowvault, empty for v2)'
+ required: false
+ type: string
+ default: ''
+
jobs:
build-and-deploy:
runs-on: ubuntu-latest
@@ -29,7 +46,7 @@ jobs:
- name: Resolve Branch for the Tagged Commit
id: resolve-branch
- if: ${{ inputs.tag == 'beta' || inputs.tag == 'public' }}
+ if: ${{ inputs.tag == 'beta' || inputs.tag == 'public' }}
run: |
TAG_COMMIT=$(git rev-list -n 1 ${{ github.ref_name }})
@@ -47,18 +64,31 @@ jobs:
id: previoustag
uses: WyriHaximus/github-action-get-previous-tag@v1
with:
- fallback: 1.0.0
+ fallback: ${{ inputs.tag-prefix }}1.0.0
+ pattern: ${{ inputs.tag-prefix }}[0-9]*.[0-9]*.[0-9]*
+
+ - name: Resolve version number
+ id: version
+ env:
+ PREVIOUS_TAG: ${{ steps.previoustag.outputs.tag }}
+ TAG_PREFIX: ${{ inputs.tag-prefix }}
+ run: |
+ TAG="$PREVIOUS_TAG"
+ VERSION="${TAG#$TAG_PREFIX}"
+ echo "version=$VERSION" >> $GITHUB_OUTPUT
- name: Bump Version
+ working-directory: ${{ inputs.variant }}
run: |
- chmod +x ./ci-scripts/bump_version.sh
+ chmod +x ../ci-scripts/bump_version.sh
if ${{ inputs.tag == 'internal' }}; then
- ./ci-scripts/bump_version.sh "${{ steps.previoustag.outputs.tag }}" "$(git rev-parse --short "$GITHUB_SHA")"
+ ../ci-scripts/bump_version.sh "${{ steps.version.outputs.version }}" "$(git rev-parse --short "$GITHUB_SHA")" "${{ inputs.package-name }}"
else
- ./ci-scripts/bump_version.sh "${{ steps.previoustag.outputs.tag }}"
+ ../ci-scripts/bump_version.sh "${{ steps.version.outputs.version }}" "" "${{ inputs.package-name }}"
fi
- name: Commit changes
+ working-directory: ${{ inputs.variant }}
run: |
git config user.name "${{ github.actor }}"
git config user.email "${{ github.actor }}@users.noreply.github.com"
@@ -68,27 +98,29 @@ jobs:
fi
git add setup.py
- git add skyflow/utils/_version.py
+ git add ${{ inputs.package-name }}/utils/_version.py
if [[ "${{ inputs.tag }}" == "internal" ]]; then
- VERSION="${{ steps.previoustag.outputs.tag }}.dev0+$(git rev-parse --short $GITHUB_SHA)"
- COMMIT_MESSAGE="[AUTOMATED] Private Release $VERSION"
+ VERSION="${{ steps.version.outputs.version }}.dev0+$(git rev-parse --short $GITHUB_SHA)"
+ COMMIT_MESSAGE="[AUTOMATED] Private Release (${{ inputs.variant }}) $VERSION"
git commit -m "$COMMIT_MESSAGE"
git push origin ${{ github.ref_name }} -f
fi
if [[ "${{ inputs.tag }}" == "beta" || "${{ inputs.tag }}" == "public" ]]; then
- COMMIT_MESSAGE="[AUTOMATED] Public Release - ${{ steps.previoustag.outputs.tag }}"
+ COMMIT_MESSAGE="[AUTOMATED] Public Release (${{ inputs.variant }}) - ${{ steps.previoustag.outputs.tag }}"
git commit -m "$COMMIT_MESSAGE"
git push origin ${{ env.branch_name }}
fi
- - name: Build and install skyflow package
+ - name: Build and install package
+ working-directory: ${{ inputs.variant }}
run: |
python setup.py sdist bdist_wheel
- pip install dist/skyflow-*.whl
+ pip install dist/*.whl
- name: Build and Publish Package
- if: ${{ inputs.tag == 'beta' || inputs.tag == 'public' }}
+ if: ${{ inputs.tag == 'beta' || inputs.tag == 'public' }}
+ working-directory: ${{ inputs.variant }}
env:
TWINE_USERNAME: __token__
TWINE_PASSWORD: ${{ secrets.PYPI_PUBLISH_TOKEN }}
@@ -98,9 +130,10 @@ jobs:
- name: Build and Publish to JFrog Artifactory
if: ${{ inputs.tag == 'internal' }}
+ working-directory: ${{ inputs.variant }}
env:
TWINE_USERNAME: ${{ secrets.JFROG_USERNAME }}
TWINE_PASSWORD: ${{ secrets.JFROG_PASSWORD }}
run: |
python setup.py sdist bdist_wheel
- twine upload --repository-url https://prekarilabs.jfrog.io/artifactory/api/pypi/skyflow-python/ dist/*
\ No newline at end of file
+ twine upload --repository-url https://prekarilabs.jfrog.io/artifactory/api/pypi/skyflow-python/ dist/*
diff --git a/.github/workflows/shared-tests.yml b/.github/workflows/shared-tests.yml
index 24fa6a0e..796eb89f 100644
--- a/.github/workflows/shared-tests.yml
+++ b/.github/workflows/shared-tests.yml
@@ -7,6 +7,20 @@ on:
description: 'Python version to use'
required: true
type: string
+ variant:
+ description: 'Build variant directory to test (v2 or flowvault)'
+ required: true
+ type: string
+ package-name:
+ description: 'Importable package name for this variant (e.g. skyflow or skyflow_flowvault)'
+ required: false
+ type: string
+ default: 'skyflow'
+ coverage-omit:
+ description: 'Comma-separated coverage --omit patterns, relative to the variant directory'
+ required: false
+ type: string
+ default: 'skyflow/generated/*'
jobs:
run-tests:
@@ -23,12 +37,14 @@ jobs:
with:
name: "credentials.json"
json: ${{ secrets.VALID_SKYFLOW_CREDS_TEST }}
+ dir: ${{ inputs.variant }}
- - name: Build and install skyflow package
+ - name: Build and install package
+ working-directory: ${{ inputs.variant }}
run: |
pip install --upgrade pip setuptools wheel
python setup.py sdist bdist_wheel
- pip install dist/skyflow-*.whl
+ pip install dist/*.whl
pip install ".[dev]"
- name: Run Spell Check
@@ -36,19 +52,21 @@ jobs:
- name: Run Linter Ruff
run: ruff check . --output-format=github
-
+
- name: 'Run Tests'
+ working-directory: ${{ inputs.variant }}
run: |
pip install -r requirements.txt
- python -m coverage run --source=skyflow --omit=skyflow/generated/*,skyflow/utils/validations/*,skyflow/vault/data/*,skyflow/vault/detect/*,skyflow/vault/tokens/*,skyflow/vault/connection/*,skyflow/error/*,skyflow/utils/enums/*,skyflow/vault/controller/_audit.py,skyflow/vault/controller/_bin_look_up.py -m unittest discover
+ python -m coverage run --source=${{ inputs.package-name }} --omit=${{ inputs.coverage-omit }} -m unittest discover
- name: coverage
+ working-directory: ${{ inputs.variant }}
run: coverage xml -o test-coverage.xml
- name: Codecov
uses: codecov/codecov-action@v2.1.0
with:
token: ${{ secrets.CODECOV_REPO_UPLOAD_TOKEN }}
- files: test-coverage.xml
- name: codecov-skyflow-python
+ files: ${{ inputs.variant }}/test-coverage.xml
+ name: codecov-skyflow-python-${{ inputs.variant }}
verbose: true
diff --git a/.gitignore b/.gitignore
index 0414a3a5..cceb4f18 100644
--- a/.gitignore
+++ b/.gitignore
@@ -104,6 +104,7 @@ celerybeat.pid
# Environments
.env
.venv
+.venv-*
env/
venv/
ENV/
diff --git a/.semgreprules/customRule.yml b/.semgreprules/customRule.yml
index b275e280..f80b5d4b 100644
--- a/.semgreprules/customRule.yml
+++ b/.semgreprules/customRule.yml
@@ -12,7 +12,7 @@ rules:
- golang
- docker
patterns:
- - pattern-regex: (?i)\b(api[_-]key|api[_-]token|api[_-]secret[_-]key|api[_-]password|token|secret[_-]key|password|auth[_-]key|auth[_-]token|AUTH_PASSWORD)\s*[:=]\s*(['"]?)((?!YOUR_EXCLUSION_PATTERN_HERE)[A-Z]+.*?)\2
+ - pattern-regex: (?i)\b(api[_-]key|api[_-]token|api[_-]secret[_-]key|api[_-]password|token|secret[_-]key|password|auth[_-]key|auth[_-]token|AUTH_PASSWORD)\s*[:=]\s*(['"])(?!\1\2)((?!YOUR_EXCLUSION_PATTERN_HERE)[A-Z]+.*?)\2
- id: check-logger-appconfig
message: >-
diff --git a/ci-scripts/bump_version.sh b/ci-scripts/bump_version.sh
index ab79e8aa..0fc6e782 100755
--- a/ci-scripts/bump_version.sh
+++ b/ci-scripts/bump_version.sh
@@ -1,13 +1,14 @@
Version=$1
SEMVER=$Version
+PackageName=${3:-skyflow}
if [ -z "$2" ]
then
echo "Bumping package version to $1"
sed -E "s/current_version = .+/current_version = '$SEMVER'/g" setup.py > tempfile && cat tempfile > setup.py && rm -f tempfile
- sed -E "s/SDK_VERSION = .+/SDK_VERSION = '$SEMVER'/g" skyflow/utils/_version.py > tempfile && cat tempfile > skyflow/utils/_version.py && rm -f tempfile
- sed -E "s/__version__ = .+/__version__ = '$SEMVER'/g" skyflow/generated/rest/version.py > tempfile && cat tempfile > skyflow/generated/rest/version.py && rm -f tempfile
+ sed -E "s/SDK_VERSION = .+/SDK_VERSION = '$SEMVER'/g" $PackageName/utils/_version.py > tempfile && cat tempfile > $PackageName/utils/_version.py && rm -f tempfile
+ sed -E "s/__version__ = .+/__version__ = '$SEMVER'/g" $PackageName/generated/rest/version.py > tempfile && cat tempfile > $PackageName/generated/rest/version.py && rm -f tempfile
echo --------------------------
echo "Done, Package now at $1"
@@ -18,8 +19,8 @@ else
echo "Bumping package version to $DEV_VERSION"
sed -E "s/current_version = .+/current_version = '$DEV_VERSION'/g" setup.py > tempfile && cat tempfile > setup.py && rm -f tempfile
- sed -E "s/SDK_VERSION = .+/SDK_VERSION = '$DEV_VERSION'/g" skyflow/utils/_version.py > tempfile && cat tempfile > skyflow/utils/_version.py && rm -f tempfile
- sed -E "s/__version__ = .+/__version__ = '$DEV_VERSION'/g" skyflow/generated/rest/version.py > tempfile && cat tempfile > skyflow/generated/rest/version.py && rm -f tempfile
+ sed -E "s/SDK_VERSION = .+/SDK_VERSION = '$DEV_VERSION'/g" $PackageName/utils/_version.py > tempfile && cat tempfile > $PackageName/utils/_version.py && rm -f tempfile
+ sed -E "s/__version__ = .+/__version__ = '$DEV_VERSION'/g" $PackageName/generated/rest/version.py > tempfile && cat tempfile > $PackageName/generated/rest/version.py && rm -f tempfile
echo --------------------------
echo "Done, Package now at $DEV_VERSION"
diff --git a/common/__init__.py b/common/__init__.py
new file mode 100644
index 00000000..984fd32b
--- /dev/null
+++ b/common/__init__.py
@@ -0,0 +1,3 @@
+# common.utils and common.errors mutually depend on each other -- forcing utils to load first
+# here avoids the circular import (mirrors skyflow/__init__.py's own first line).
+from . import utils # noqa: F401
diff --git a/skyflow/generated/__init__.py b/common/client/__init__.py
similarity index 100%
rename from skyflow/generated/__init__.py
rename to common/client/__init__.py
diff --git a/common/client/base_skyflow.py b/common/client/base_skyflow.py
new file mode 100644
index 00000000..29cd1a7c
--- /dev/null
+++ b/common/client/base_skyflow.py
@@ -0,0 +1,341 @@
+from abc import ABC, abstractmethod
+from collections import OrderedDict
+from common.errors import SkyflowError
+from common.utils import SkyflowMessages
+from common.utils.logger import log_info, log_warn
+from common.utils.constants import OptionField
+
+_BUILDER_TEMPLATE_ERROR = (
+ "BaseSkyflowImpl.Builder is an interface template -- build a concrete Skyflow "
+ "class via make_skyflow_class() instead of using it directly. Missing: {missing}"
+)
+_CONNECTIONS_NOT_SUPPORTED_ERROR = "Connections are not supported by this Skyflow SDK variant"
+_DETECT_NOT_SUPPORTED_ERROR = "Detect is not supported by this Skyflow SDK variant"
+
+
+class BaseSkyflow(ABC):
+ @classmethod
+ @abstractmethod
+ def builder(cls):
+ raise NotImplementedError
+
+ @abstractmethod
+ def add_vault_config(self, config):
+ raise NotImplementedError
+
+ @abstractmethod
+ def remove_vault_config(self, vault_id):
+ raise NotImplementedError
+
+ @abstractmethod
+ def update_vault_config(self, config):
+ raise NotImplementedError
+
+ @abstractmethod
+ def get_vault_config(self, vault_id):
+ raise NotImplementedError
+
+ @abstractmethod
+ def add_skyflow_credentials(self, credentials):
+ raise NotImplementedError
+
+ @abstractmethod
+ def update_skyflow_credentials(self, credentials):
+ raise NotImplementedError
+
+ @abstractmethod
+ def set_log_level(self, log_level):
+ raise NotImplementedError
+
+ @abstractmethod
+ def update_log_level(self, log_level):
+ raise NotImplementedError
+
+ @abstractmethod
+ def get_log_level(self):
+ raise NotImplementedError
+
+ @abstractmethod
+ def vault(self, vault_id=None):
+ raise NotImplementedError
+
+
+class BaseSkyflowImpl(BaseSkyflow):
+
+ def __init__(self, builder):
+ if type(self) is BaseSkyflowImpl:
+ raise SkyflowError(
+ SkyflowMessages.Error.BASE_SKYFLOW_INSTANTIATION_NOT_ALLOWED.value,
+ SkyflowMessages.ErrorCodes.INVALID_INPUT.value,
+ )
+ self.__builder = builder
+ log_info(self.__builder._skyflow_messages.Info.CLIENT_INITIALIZED.value, self.__builder.get_logger())
+
+ @classmethod
+ def builder(cls):
+ return cls.Builder()
+
+ def add_vault_config(self, config):
+ self.__builder._add_vault_config(config)
+ return self
+
+ def remove_vault_config(self, vault_id):
+ self.__builder.remove_vault_config(vault_id)
+
+ def update_vault_config(self, config):
+ self.__builder.update_vault_config(config)
+
+ def get_vault_config(self, vault_id):
+ return self.__builder.get_vault_config(vault_id).get(OptionField.VAULT_CLIENT).get_config()
+
+ def add_skyflow_credentials(self, credentials):
+ self.__builder._add_skyflow_credentials(credentials)
+ return self
+
+ def update_skyflow_credentials(self, credentials):
+ self.__builder._add_skyflow_credentials(credentials)
+
+ def set_log_level(self, log_level):
+ self.__builder._set_log_level(log_level)
+ return self
+
+ def update_log_level(self, log_level):
+ """.. deprecated:: Use set_log_level() instead. Will be removed in a future release."""
+ log_warn(self.__builder._skyflow_messages.Warning.UPDATE_LOG_LEVEL_DEPRECATED.value)
+ return self.set_log_level(log_level)
+
+ def get_log_level(self):
+ return self.__builder.get_log_level()
+
+ def vault(self, vault_id=None):
+ vault_config = self.__builder.get_vault_config(vault_id)
+ return vault_config.get(OptionField.VAULT_CONTROLLER)
+
+ def _get_builder(self):
+ return self.__builder
+
+ class Builder(ABC):
+ _vault_client_cls = None
+ _vault_controller_cls = None
+ _connection_cls = None
+ _detect_cls = None
+ _logger_cls = None
+ _default_log_level = None
+ _skyflow_messages = None
+ _skyflow_cls = None
+ _validate_vault_config = None
+ _validate_update_vault_config = None
+ _validate_connection_config = None
+ _validate_update_connection_config = None
+ _validate_log_level = None
+ _validate_credentials = None
+ _set_active_log_level = None
+
+ _REQUIRED_HOOKS = (
+ '_vault_client_cls', '_vault_controller_cls', '_logger_cls', '_default_log_level',
+ '_skyflow_messages', '_skyflow_cls', '_validate_vault_config',
+ '_validate_update_vault_config', '_validate_log_level', '_validate_credentials',
+ )
+
+ def __init__(self):
+ missing = [hook for hook in self._REQUIRED_HOOKS if getattr(self, hook) is None]
+ if missing:
+ raise NotImplementedError(_BUILDER_TEMPLATE_ERROR.format(missing=', '.join(missing)))
+ self.__vault_configs = OrderedDict()
+ self.__vault_list = list()
+ self.__connection_configs = OrderedDict()
+ self.__connection_list = list()
+ self.__skyflow_credentials = None
+ self.__log_level = self._default_log_level
+ self.__logger = self._logger_cls(self._default_log_level)
+
+ def _require_connections(self):
+ if self._connection_cls is None:
+ raise NotImplementedError(_CONNECTIONS_NOT_SUPPORTED_ERROR)
+
+ def _require_detect(self):
+ if self._detect_cls is None:
+ raise NotImplementedError(_DETECT_NOT_SUPPORTED_ERROR)
+
+ def add_vault_config(self, config):
+ vault_id = config.get(OptionField.VAULT_ID)
+ if not isinstance(vault_id, str) or not vault_id:
+ raise SkyflowError(
+ self._skyflow_messages.Error.INVALID_VAULT_ID.value,
+ self._skyflow_messages.ErrorCodes.INVALID_INPUT.value
+ )
+ if vault_id in [vault.get(OptionField.VAULT_ID) for vault in self.__vault_list]:
+ log_info(self._skyflow_messages.Info.VAULT_CONFIG_EXISTS.value.format(vault_id), self.__logger)
+ raise SkyflowError(
+ self._skyflow_messages.Error.VAULT_ID_ALREADY_EXISTS.value.format(vault_id),
+ self._skyflow_messages.ErrorCodes.INVALID_INPUT.value
+ )
+ self.__vault_list.append(config)
+ return self
+
+ def remove_vault_config(self, vault_id):
+ if vault_id in self.__vault_configs.keys():
+ self.__vault_configs.pop(vault_id)
+ else:
+ raise SkyflowError(self._skyflow_messages.Error.INVALID_VAULT_ID.value,
+ self._skyflow_messages.ErrorCodes.INVALID_INPUT.value)
+
+ def update_vault_config(self, config):
+ self._validate_update_vault_config(self.__logger, config)
+ vault_id = config.get(OptionField.VAULT_ID)
+ if vault_id not in self.__vault_configs:
+ raise SkyflowError(self._skyflow_messages.Error.VAULT_ID_NOT_IN_CONFIG_LIST.value.format(vault_id), self._skyflow_messages.ErrorCodes.INVALID_INPUT.value)
+ vault_config = self.__vault_configs[vault_id]
+ vault_config.get(OptionField.VAULT_CLIENT).update_config(config)
+
+ def get_vault_config(self, vault_id):
+ if vault_id is None:
+ if self.__vault_configs:
+ return next(iter(self.__vault_configs.values()))
+ raise SkyflowError(self._skyflow_messages.Error.EMPTY_VAULT_CONFIGS.value, self._skyflow_messages.ErrorCodes.INVALID_INPUT.value)
+
+ if vault_id in self.__vault_configs:
+ return self.__vault_configs.get(vault_id)
+ log_info(self._skyflow_messages.Info.VAULT_CONFIG_DOES_NOT_EXIST.value.format(vault_id), self.__logger)
+ raise SkyflowError(self._skyflow_messages.Error.VAULT_ID_NOT_IN_CONFIG_LIST.value.format(vault_id), self._skyflow_messages.ErrorCodes.INVALID_INPUT.value)
+
+ def add_connection_config(self, config):
+ self._require_connections()
+ connection_id = config.get(OptionField.CONNECTION_ID)
+ if not isinstance(connection_id, str) or not connection_id:
+ raise SkyflowError(
+ self._skyflow_messages.Error.INVALID_CONNECTION_ID.value,
+ self._skyflow_messages.ErrorCodes.INVALID_INPUT.value
+ )
+ if connection_id in [connection.get(OptionField.CONNECTION_ID) for connection in self.__connection_list]:
+ log_info(self._skyflow_messages.Info.CONNECTION_CONFIG_EXISTS.value.format(connection_id), self.__logger)
+ raise SkyflowError(
+ self._skyflow_messages.Error.CONNECTION_ID_ALREADY_EXISTS.value.format(connection_id),
+ self._skyflow_messages.ErrorCodes.INVALID_INPUT.value
+ )
+ self.__connection_list.append(config)
+ return self
+
+ def remove_connection_config(self, connection_id):
+ self._require_connections()
+ if connection_id in self.__connection_configs.keys():
+ self.__connection_configs.pop(connection_id)
+ else:
+ raise SkyflowError(self._skyflow_messages.Error.INVALID_CONNECTION_ID.value,
+ self._skyflow_messages.ErrorCodes.INVALID_INPUT.value)
+
+ def update_connection_config(self, config):
+ self._require_connections()
+ self._validate_update_connection_config(self.__logger, config)
+ connection_id = config.get(OptionField.CONNECTION_ID)
+ if connection_id not in self.__connection_configs:
+ raise SkyflowError(self._skyflow_messages.Error.CONNECTION_ID_NOT_IN_CONFIG_LIST.value.format(connection_id), self._skyflow_messages.ErrorCodes.INVALID_INPUT.value)
+ connection_config = self.__connection_configs[connection_id]
+ connection_config.get(OptionField.VAULT_CLIENT).update_config(config)
+
+ def get_connection_config(self, connection_id):
+ self._require_connections()
+ if connection_id is None:
+ if self.__connection_configs:
+ return next(iter(self.__connection_configs.values()))
+
+ raise SkyflowError(self._skyflow_messages.Error.EMPTY_CONNECTION_CONFIGS.value, self._skyflow_messages.ErrorCodes.INVALID_INPUT.value)
+
+ if connection_id in self.__connection_configs:
+ return self.__connection_configs.get(connection_id)
+ log_info(self._skyflow_messages.Info.CONNECTION_CONFIG_DOES_NOT_EXIST.value.format(connection_id), self.__logger)
+ raise SkyflowError(self._skyflow_messages.Error.CONNECTION_ID_NOT_IN_CONFIG_LIST.value.format(connection_id), self._skyflow_messages.ErrorCodes.INVALID_INPUT.value)
+
+ def add_skyflow_credentials(self, credentials):
+ self.__skyflow_credentials = credentials
+ return self
+
+ def set_log_level(self, log_level):
+ self.__log_level = log_level
+ return self
+
+ def get_logger(self):
+ return self.__logger
+
+ def get_log_level(self):
+ return self.__log_level
+
+ def _add_vault_config(self, config):
+ self._validate_vault_config(self.__logger, config)
+ vault_id = config.get(OptionField.VAULT_ID)
+ if vault_id in self.__vault_configs:
+ raise SkyflowError(
+ self._skyflow_messages.Error.VAULT_ID_ALREADY_EXISTS.value.format(vault_id),
+ self._skyflow_messages.ErrorCodes.INVALID_INPUT.value
+ )
+ vault_client = self._vault_client_cls(config)
+ vault_config = {
+ OptionField.VAULT_CLIENT: vault_client,
+ OptionField.VAULT_CONTROLLER: self._vault_controller_cls(vault_client),
+ }
+ if self._detect_cls is not None:
+ vault_config[OptionField.DETECT_CONTROLLER] = self._detect_cls(vault_client)
+ self.__vault_configs[vault_id] = vault_config
+ log_info(self._skyflow_messages.Info.VAULT_CONTROLLER_INITIALIZED.value.format(vault_id), self.__logger)
+ if self._detect_cls is not None:
+ log_info(self._skyflow_messages.Info.DETECT_CONTROLLER_INITIALIZED.value.format(vault_id), self.__logger)
+
+ def _add_connection_config(self, config):
+ self._validate_connection_config(self.__logger, config)
+ connection_id = config.get(OptionField.CONNECTION_ID)
+ if connection_id in self.__connection_configs:
+ raise SkyflowError(
+ self._skyflow_messages.Error.CONNECTION_ID_ALREADY_EXISTS.value.format(connection_id),
+ self._skyflow_messages.ErrorCodes.INVALID_INPUT.value
+ )
+ vault_client = self._vault_client_cls(config)
+ self.__connection_configs[connection_id] = {
+ OptionField.VAULT_CLIENT: vault_client,
+ OptionField.CONTROLLER: self._connection_cls(vault_client)
+ }
+ log_info(self._skyflow_messages.Info.CONNECTION_CONTROLLER_INITIALIZED.value.format(connection_id), self.__logger)
+
+ def _update_vault_client_logger(self, log_level, logger):
+ for vault_id, vault_config in self.__vault_configs.items():
+ vault_config.get(OptionField.VAULT_CLIENT).set_logger(log_level, logger)
+
+ for connection_id, connection_config in self.__connection_configs.items():
+ connection_config.get(OptionField.VAULT_CLIENT).set_logger(log_level, logger)
+
+ def _set_log_level(self, log_level):
+ self._validate_log_level(self.__logger, log_level)
+ self.__log_level = log_level
+ self.__logger.set_log_level(log_level)
+ if self._set_active_log_level is not None:
+ self._set_active_log_level(log_level)
+ self._update_vault_client_logger(log_level, self.__logger)
+ log_info(self._skyflow_messages.Info.LOGGER_SETUP_DONE.value, self.__logger)
+ log_info(self._skyflow_messages.Info.CURRENT_LOG_LEVEL.value.format(self.__log_level), self.__logger)
+
+ def _add_skyflow_credentials(self, credentials):
+ if credentials is not None:
+ self.__skyflow_credentials = credentials
+ self._validate_credentials(self.__logger, credentials)
+ for vault_id, vault_config in self.__vault_configs.items():
+ vault_config.get(OptionField.VAULT_CLIENT).set_common_skyflow_credentials(credentials)
+
+ for connection_id, connection_config in self.__connection_configs.items():
+ connection_config.get(OptionField.VAULT_CLIENT).set_common_skyflow_credentials(self.__skyflow_credentials)
+
+ def build(self):
+ self._validate_log_level(self.__logger, self.__log_level)
+ self.__logger.set_log_level(self.__log_level)
+ if self._set_active_log_level is not None:
+ self._set_active_log_level(self.__log_level)
+
+ for config in self.__vault_list:
+ self._add_vault_config(config)
+
+ for config in self.__connection_list:
+ self._add_connection_config(config)
+
+ self._update_vault_client_logger(self.__log_level, self.__logger)
+
+ self._add_skyflow_credentials(self.__skyflow_credentials)
+
+ return self._skyflow_cls(self)
diff --git a/common/client/utils/__init__.py b/common/client/utils/__init__.py
new file mode 100644
index 00000000..5c114c7f
--- /dev/null
+++ b/common/client/utils/__init__.py
@@ -0,0 +1 @@
+from common.client.utils._utils import ConnectionCapable, DetectCapable, ConnectionMixin, DetectMixin, make_skyflow_class
diff --git a/common/client/utils/_utils.py b/common/client/utils/_utils.py
new file mode 100644
index 00000000..5815dd70
--- /dev/null
+++ b/common/client/utils/_utils.py
@@ -0,0 +1,127 @@
+from abc import ABC, abstractmethod
+from functools import partial
+
+from common.utils.constants import OptionField
+from common.utils.enums import LogLevel as _CommonLogLevel
+from common.utils.logger import Logger as _CommonLogger
+from common.utils.validations import (
+ validate_vault_config as _common_validate_vault_config,
+ validate_update_vault_config as _common_validate_update_vault_config,
+ validate_log_level as _common_validate_log_level,
+ validate_credentials as _common_validate_credentials,
+)
+from common.client.base_skyflow import BaseSkyflowImpl
+
+
+class ConnectionCapable(ABC):
+ @abstractmethod
+ def add_connection_config(self, config):
+ raise NotImplementedError
+
+ @abstractmethod
+ def remove_connection_config(self, connection_id):
+ raise NotImplementedError
+
+ @abstractmethod
+ def update_connection_config(self, config):
+ raise NotImplementedError
+
+ @abstractmethod
+ def get_connection_config(self, connection_id):
+ raise NotImplementedError
+
+ @abstractmethod
+ def connection(self, connection_id=None):
+ raise NotImplementedError
+
+
+class DetectCapable(ABC):
+ @abstractmethod
+ def detect(self, vault_id=None):
+ raise NotImplementedError
+
+
+class ConnectionMixin(ConnectionCapable):
+
+ def add_connection_config(self, config):
+ builder = self._get_builder()
+ builder._require_connections()
+ builder._add_connection_config(config)
+ return self
+
+ def remove_connection_config(self, connection_id):
+ builder = self._get_builder()
+ builder._require_connections()
+ builder.remove_connection_config(connection_id)
+ return self
+
+ def update_connection_config(self, config):
+ builder = self._get_builder()
+ builder._require_connections()
+ builder.update_connection_config(config)
+ return self
+
+ def get_connection_config(self, connection_id):
+ builder = self._get_builder()
+ builder._require_connections()
+ return builder.get_connection_config(connection_id).get(OptionField.VAULT_CLIENT).get_config()
+
+ def connection(self, connection_id=None):
+ builder = self._get_builder()
+ builder._require_connections()
+ connection_config = builder.get_connection_config(connection_id)
+ return connection_config.get(OptionField.CONTROLLER)
+
+
+class DetectMixin(DetectCapable):
+
+ def detect(self, vault_id=None):
+ builder = self._get_builder()
+ builder._require_detect()
+ vault_config = builder.get_vault_config(vault_id)
+ return vault_config.get(OptionField.DETECT_CONTROLLER)
+
+
+def make_skyflow_class(*, vault_client_cls, vault_controller_cls, skyflow_messages,
+ validate_vault_config=None, validate_update_vault_config=None,
+ validate_log_level=None, validate_credentials=None,
+ logger_cls=_CommonLogger, default_log_level=_CommonLogLevel.ERROR,
+ connection_cls=None, detect_cls=None,
+ validate_connection_config=None, validate_update_connection_config=None,
+ set_active_log_level=None):
+
+ if connection_cls is not None and (validate_connection_config is None or validate_update_connection_config is None):
+ raise ValueError("connection_cls requires validate_connection_config and validate_update_connection_config")
+
+ validate_vault_config = validate_vault_config or partial(_common_validate_vault_config, messages=skyflow_messages)
+ validate_update_vault_config = validate_update_vault_config or partial(_common_validate_update_vault_config, messages=skyflow_messages)
+ validate_log_level = validate_log_level or partial(_common_validate_log_level, messages=skyflow_messages)
+ validate_credentials = validate_credentials or partial(_common_validate_credentials, messages=skyflow_messages)
+
+ builder_attrs = {
+ '_vault_client_cls': vault_client_cls,
+ '_vault_controller_cls': vault_controller_cls,
+ '_connection_cls': connection_cls,
+ '_detect_cls': detect_cls,
+ '_logger_cls': logger_cls,
+ '_default_log_level': default_log_level,
+ '_skyflow_messages': skyflow_messages,
+ '_validate_vault_config': staticmethod(validate_vault_config),
+ '_validate_update_vault_config': staticmethod(validate_update_vault_config),
+ '_validate_connection_config': staticmethod(validate_connection_config) if validate_connection_config else None,
+ '_validate_update_connection_config': staticmethod(validate_update_connection_config) if validate_update_connection_config else None,
+ '_validate_log_level': staticmethod(validate_log_level),
+ '_validate_credentials': staticmethod(validate_credentials),
+ '_set_active_log_level': staticmethod(set_active_log_level) if set_active_log_level else None,
+ }
+ variant_builder = type('Builder', (BaseSkyflowImpl.Builder,), builder_attrs)
+
+ bases = [BaseSkyflowImpl]
+ if connection_cls is not None:
+ bases.append(ConnectionMixin)
+ if detect_cls is not None:
+ bases.append(DetectMixin)
+
+ variant_skyflow = type('Skyflow', tuple(bases), {'Builder': variant_builder})
+ variant_builder._skyflow_cls = variant_skyflow
+ return variant_skyflow
diff --git a/common/errors/__init__.py b/common/errors/__init__.py
new file mode 100644
index 00000000..17a2fe39
--- /dev/null
+++ b/common/errors/__init__.py
@@ -0,0 +1 @@
+from ._skyflow_error import SkyflowError
diff --git a/skyflow/error/_skyflow_error.py b/common/errors/_skyflow_error.py
similarity index 93%
rename from skyflow/error/_skyflow_error.py
rename to common/errors/_skyflow_error.py
index cda064e2..8ee07694 100644
--- a/skyflow/error/_skyflow_error.py
+++ b/common/errors/_skyflow_error.py
@@ -1,4 +1,4 @@
-from skyflow.utils import SkyflowMessages
+from common.utils import SkyflowMessages
class SkyflowError(Exception):
def __init__(self,
diff --git a/skyflow/service_account/client/__init__.py b/common/generated/__init__.py
similarity index 100%
rename from skyflow/service_account/client/__init__.py
rename to common/generated/__init__.py
diff --git a/common/generated/requirements.txt b/common/generated/requirements.txt
new file mode 100644
index 00000000..e80f640a
--- /dev/null
+++ b/common/generated/requirements.txt
@@ -0,0 +1,4 @@
+httpx>=0.21.2
+pydantic>= 1.9.2
+pydantic-core>=2.18.2
+typing_extensions>= 4.0.0
diff --git a/common/generated/rest/__init__.py b/common/generated/rest/__init__.py
new file mode 100644
index 00000000..2b7e1ccf
--- /dev/null
+++ b/common/generated/rest/__init__.py
@@ -0,0 +1,24 @@
+# This file was auto-generated by Fern from our API Definition.
+
+# isort: skip_file
+
+from .types import GooglerpcStatus, ProtobufAny, V1GetAuthTokenResponse
+from .errors import BadRequestError, NotFoundError, UnauthorizedError
+from . import authentication
+from .client import AsyncSkyflowAuth, SkyflowAuth
+from .environment import SkyflowAuthEnvironment
+from .version import __version__
+
+__all__ = [
+ "AsyncSkyflowAuth",
+ "BadRequestError",
+ "GooglerpcStatus",
+ "NotFoundError",
+ "ProtobufAny",
+ "SkyflowAuth",
+ "SkyflowAuthEnvironment",
+ "UnauthorizedError",
+ "V1GetAuthTokenResponse",
+ "__version__",
+ "authentication",
+]
diff --git a/skyflow/generated/rest/authentication/__init__.py b/common/generated/rest/authentication/__init__.py
similarity index 100%
rename from skyflow/generated/rest/authentication/__init__.py
rename to common/generated/rest/authentication/__init__.py
diff --git a/common/generated/rest/authentication/client.py b/common/generated/rest/authentication/client.py
new file mode 100644
index 00000000..9653b6d3
--- /dev/null
+++ b/common/generated/rest/authentication/client.py
@@ -0,0 +1,181 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
+from ..core.request_options import RequestOptions
+from ..types.v_1_get_auth_token_response import V1GetAuthTokenResponse
+from .raw_client import AsyncRawAuthenticationClient, RawAuthenticationClient
+
+# this is used as the default value for optional parameters
+OMIT = typing.cast(typing.Any, ...)
+
+
+class AuthenticationClient:
+ def __init__(self, *, client_wrapper: SyncClientWrapper):
+ self._raw_client = RawAuthenticationClient(client_wrapper=client_wrapper)
+
+ @property
+ def with_raw_response(self) -> RawAuthenticationClient:
+ """
+ Retrieves a raw implementation of this client that returns raw responses.
+
+ Returns
+ -------
+ RawAuthenticationClient
+ """
+ return self._raw_client
+
+ def authentication_service_get_auth_token(
+ self,
+ *,
+ grant_type: str,
+ assertion: str,
+ subject_token: typing.Optional[str] = OMIT,
+ subject_token_type: typing.Optional[str] = OMIT,
+ requested_token_use: typing.Optional[str] = OMIT,
+ scope: typing.Optional[str] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1GetAuthTokenResponse:
+ """
+ Generates a Bearer Token to authenticate with Skyflow. This method doesn't require the Authorization header.
Note: For recommended ways to authenticate, see API authentication.
+
+ Parameters
+ ----------
+ grant_type : str
+ Grant type of the request. Set this to `urn:ietf:params:oauth:grant-type:jwt-bearer`.
+
+ assertion : str
+ User-signed JWT token that contains the following fields:
iss: Issuer of the JWT.key: Unique identifier for the key.aud: Recipient the JWT is intended for.exp: Time the JWT expires.sub: Subject of the JWT.ctx: (Optional) Value for Context-aware authorization.
+
+ subject_token : typing.Optional[str]
+ Subject token.
+
+ subject_token_type : typing.Optional[str]
+ Subject token type.
+
+ requested_token_use : typing.Optional[str]
+ Token use type. Either `delegation` or `impersonation`.
+
+ scope : typing.Optional[str]
+ Subset of available roles to associate with the requested token. Uses the format "role:\ role:\".
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1GetAuthTokenResponse
+ A successful response.
+
+ Examples
+ --------
+ from skyflow import SkyflowAuth
+
+ client = SkyflowAuth(
+ token="YOUR_TOKEN",
+ )
+ client.authentication.authentication_service_get_auth_token(
+ grant_type="urn:ietf:params:oauth:grant-type:jwt-bearer",
+ assertion="eyLhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaXNzIjoiY29tcGFueSIsImV4cCI6MTYxNTE5MzgwNywiaWF0IjoxNjE1MTY1MDQwLCJhdWQiOiKzb21lYXVkaWVuY2UifQ.4pcPyMDQ9o1PSyXnrXCjTwXyr4BSezdI1AVTmud2fU3",
+ )
+ """
+ _response = self._raw_client.authentication_service_get_auth_token(
+ grant_type=grant_type,
+ assertion=assertion,
+ subject_token=subject_token,
+ subject_token_type=subject_token_type,
+ requested_token_use=requested_token_use,
+ scope=scope,
+ request_options=request_options,
+ )
+ return _response.data
+
+
+class AsyncAuthenticationClient:
+ def __init__(self, *, client_wrapper: AsyncClientWrapper):
+ self._raw_client = AsyncRawAuthenticationClient(client_wrapper=client_wrapper)
+
+ @property
+ def with_raw_response(self) -> AsyncRawAuthenticationClient:
+ """
+ Retrieves a raw implementation of this client that returns raw responses.
+
+ Returns
+ -------
+ AsyncRawAuthenticationClient
+ """
+ return self._raw_client
+
+ async def authentication_service_get_auth_token(
+ self,
+ *,
+ grant_type: str,
+ assertion: str,
+ subject_token: typing.Optional[str] = OMIT,
+ subject_token_type: typing.Optional[str] = OMIT,
+ requested_token_use: typing.Optional[str] = OMIT,
+ scope: typing.Optional[str] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1GetAuthTokenResponse:
+ """
+ Generates a Bearer Token to authenticate with Skyflow. This method doesn't require the Authorization header.
Note: For recommended ways to authenticate, see API authentication.
+
+ Parameters
+ ----------
+ grant_type : str
+ Grant type of the request. Set this to `urn:ietf:params:oauth:grant-type:jwt-bearer`.
+
+ assertion : str
+ User-signed JWT token that contains the following fields:
iss: Issuer of the JWT.key: Unique identifier for the key.aud: Recipient the JWT is intended for.exp: Time the JWT expires.sub: Subject of the JWT.ctx: (Optional) Value for Context-aware authorization.
+
+ subject_token : typing.Optional[str]
+ Subject token.
+
+ subject_token_type : typing.Optional[str]
+ Subject token type.
+
+ requested_token_use : typing.Optional[str]
+ Token use type. Either `delegation` or `impersonation`.
+
+ scope : typing.Optional[str]
+ Subset of available roles to associate with the requested token. Uses the format "role:\ role:\".
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1GetAuthTokenResponse
+ A successful response.
+
+ Examples
+ --------
+ import asyncio
+
+ from skyflow import AsyncSkyflowAuth
+
+ client = AsyncSkyflowAuth(
+ token="YOUR_TOKEN",
+ )
+
+
+ async def main() -> None:
+ await client.authentication.authentication_service_get_auth_token(
+ grant_type="urn:ietf:params:oauth:grant-type:jwt-bearer",
+ assertion="eyLhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaXNzIjoiY29tcGFueSIsImV4cCI6MTYxNTE5MzgwNywiaWF0IjoxNjE1MTY1MDQwLCJhdWQiOiKzb21lYXVkaWVuY2UifQ.4pcPyMDQ9o1PSyXnrXCjTwXyr4BSezdI1AVTmud2fU3",
+ )
+
+
+ asyncio.run(main())
+ """
+ _response = await self._raw_client.authentication_service_get_auth_token(
+ grant_type=grant_type,
+ assertion=assertion,
+ subject_token=subject_token,
+ subject_token_type=subject_token_type,
+ requested_token_use=requested_token_use,
+ scope=scope,
+ request_options=request_options,
+ )
+ return _response.data
diff --git a/common/generated/rest/authentication/raw_client.py b/common/generated/rest/authentication/raw_client.py
new file mode 100644
index 00000000..bb1c2ed7
--- /dev/null
+++ b/common/generated/rest/authentication/raw_client.py
@@ -0,0 +1,241 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+from json.decoder import JSONDecodeError
+
+from ..core.api_error import ApiError
+from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
+from ..core.http_response import AsyncHttpResponse, HttpResponse
+from ..core.pydantic_utilities import parse_obj_as
+from ..core.request_options import RequestOptions
+from ..errors.bad_request_error import BadRequestError
+from ..errors.not_found_error import NotFoundError
+from ..errors.unauthorized_error import UnauthorizedError
+from ..types.v_1_get_auth_token_response import V1GetAuthTokenResponse
+
+# this is used as the default value for optional parameters
+OMIT = typing.cast(typing.Any, ...)
+
+
+class RawAuthenticationClient:
+ def __init__(self, *, client_wrapper: SyncClientWrapper):
+ self._client_wrapper = client_wrapper
+
+ def authentication_service_get_auth_token(
+ self,
+ *,
+ grant_type: str,
+ assertion: str,
+ subject_token: typing.Optional[str] = OMIT,
+ subject_token_type: typing.Optional[str] = OMIT,
+ requested_token_use: typing.Optional[str] = OMIT,
+ scope: typing.Optional[str] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> HttpResponse[V1GetAuthTokenResponse]:
+ """
+ Generates a Bearer Token to authenticate with Skyflow. This method doesn't require the Authorization header.
Note: For recommended ways to authenticate, see API authentication.
+
+ Parameters
+ ----------
+ grant_type : str
+ Grant type of the request. Set this to `urn:ietf:params:oauth:grant-type:jwt-bearer`.
+
+ assertion : str
+ User-signed JWT token that contains the following fields:
iss: Issuer of the JWT.key: Unique identifier for the key.aud: Recipient the JWT is intended for.exp: Time the JWT expires.sub: Subject of the JWT.ctx: (Optional) Value for Context-aware authorization.
+
+ subject_token : typing.Optional[str]
+ Subject token.
+
+ subject_token_type : typing.Optional[str]
+ Subject token type.
+
+ requested_token_use : typing.Optional[str]
+ Token use type. Either `delegation` or `impersonation`.
+
+ scope : typing.Optional[str]
+ Subset of available roles to associate with the requested token. Uses the format "role:\ role:\".
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ HttpResponse[V1GetAuthTokenResponse]
+ A successful response.
+ """
+ _response = self._client_wrapper.httpx_client.request(
+ "v1/auth/sa/oauth/token",
+ method="POST",
+ json={
+ "grant_type": grant_type,
+ "assertion": assertion,
+ "subject_token": subject_token,
+ "subject_token_type": subject_token_type,
+ "requested_token_use": requested_token_use,
+ "scope": scope,
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1GetAuthTokenResponse,
+ parse_obj_as(
+ type_=V1GetAuthTokenResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return HttpResponse(response=_response, data=_data)
+ if _response.status_code == 400:
+ raise BadRequestError(
+ headers=dict(_response.headers),
+ body=typing.cast(
+ typing.Dict[str, typing.Optional[typing.Any]],
+ parse_obj_as(
+ type_=typing.Dict[str, typing.Optional[typing.Any]], # type: ignore
+ object_=_response.json(),
+ ),
+ ),
+ )
+ if _response.status_code == 401:
+ raise UnauthorizedError(
+ headers=dict(_response.headers),
+ body=typing.cast(
+ typing.Dict[str, typing.Optional[typing.Any]],
+ parse_obj_as(
+ type_=typing.Dict[str, typing.Optional[typing.Any]], # type: ignore
+ object_=_response.json(),
+ ),
+ ),
+ )
+ if _response.status_code == 404:
+ raise NotFoundError(
+ headers=dict(_response.headers),
+ body=typing.cast(
+ typing.Dict[str, typing.Optional[typing.Any]],
+ parse_obj_as(
+ type_=typing.Dict[str, typing.Optional[typing.Any]], # type: ignore
+ object_=_response.json(),
+ ),
+ ),
+ )
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+
+class AsyncRawAuthenticationClient:
+ def __init__(self, *, client_wrapper: AsyncClientWrapper):
+ self._client_wrapper = client_wrapper
+
+ async def authentication_service_get_auth_token(
+ self,
+ *,
+ grant_type: str,
+ assertion: str,
+ subject_token: typing.Optional[str] = OMIT,
+ subject_token_type: typing.Optional[str] = OMIT,
+ requested_token_use: typing.Optional[str] = OMIT,
+ scope: typing.Optional[str] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> AsyncHttpResponse[V1GetAuthTokenResponse]:
+ """
+ Generates a Bearer Token to authenticate with Skyflow. This method doesn't require the Authorization header.
Note: For recommended ways to authenticate, see API authentication.
+
+ Parameters
+ ----------
+ grant_type : str
+ Grant type of the request. Set this to `urn:ietf:params:oauth:grant-type:jwt-bearer`.
+
+ assertion : str
+ User-signed JWT token that contains the following fields:
iss: Issuer of the JWT.key: Unique identifier for the key.aud: Recipient the JWT is intended for.exp: Time the JWT expires.sub: Subject of the JWT.ctx: (Optional) Value for Context-aware authorization.
+
+ subject_token : typing.Optional[str]
+ Subject token.
+
+ subject_token_type : typing.Optional[str]
+ Subject token type.
+
+ requested_token_use : typing.Optional[str]
+ Token use type. Either `delegation` or `impersonation`.
+
+ scope : typing.Optional[str]
+ Subset of available roles to associate with the requested token. Uses the format "role:\ role:\".
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ AsyncHttpResponse[V1GetAuthTokenResponse]
+ A successful response.
+ """
+ _response = await self._client_wrapper.httpx_client.request(
+ "v1/auth/sa/oauth/token",
+ method="POST",
+ json={
+ "grant_type": grant_type,
+ "assertion": assertion,
+ "subject_token": subject_token,
+ "subject_token_type": subject_token_type,
+ "requested_token_use": requested_token_use,
+ "scope": scope,
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1GetAuthTokenResponse,
+ parse_obj_as(
+ type_=V1GetAuthTokenResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return AsyncHttpResponse(response=_response, data=_data)
+ if _response.status_code == 400:
+ raise BadRequestError(
+ headers=dict(_response.headers),
+ body=typing.cast(
+ typing.Dict[str, typing.Optional[typing.Any]],
+ parse_obj_as(
+ type_=typing.Dict[str, typing.Optional[typing.Any]], # type: ignore
+ object_=_response.json(),
+ ),
+ ),
+ )
+ if _response.status_code == 401:
+ raise UnauthorizedError(
+ headers=dict(_response.headers),
+ body=typing.cast(
+ typing.Dict[str, typing.Optional[typing.Any]],
+ parse_obj_as(
+ type_=typing.Dict[str, typing.Optional[typing.Any]], # type: ignore
+ object_=_response.json(),
+ ),
+ ),
+ )
+ if _response.status_code == 404:
+ raise NotFoundError(
+ headers=dict(_response.headers),
+ body=typing.cast(
+ typing.Dict[str, typing.Optional[typing.Any]],
+ parse_obj_as(
+ type_=typing.Dict[str, typing.Optional[typing.Any]], # type: ignore
+ object_=_response.json(),
+ ),
+ ),
+ )
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
diff --git a/common/generated/rest/client.py b/common/generated/rest/client.py
new file mode 100644
index 00000000..b210e4da
--- /dev/null
+++ b/common/generated/rest/client.py
@@ -0,0 +1,153 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import httpx
+from .authentication.client import AsyncAuthenticationClient, AuthenticationClient
+from .core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
+from .environment import SkyflowAuthEnvironment
+
+
+class SkyflowAuth:
+ """
+ Use this class to access the different functions within the SDK. You can instantiate any number of clients with different configuration that will propagate to these functions.
+
+ Parameters
+ ----------
+ base_url : typing.Optional[str]
+ The base url to use for requests from the client.
+
+ environment : SkyflowAuthEnvironment
+ The environment to use for requests from the client. from .environment import SkyflowAuthEnvironment
+
+
+
+ Defaults to SkyflowAuthEnvironment.PRODUCTION
+
+
+
+ token : typing.Optional[typing.Union[str, typing.Callable[[], str]]]
+ headers : typing.Optional[typing.Dict[str, str]]
+ Additional headers to send with every request.
+
+ timeout : typing.Optional[float]
+ The timeout to be used, in seconds, for requests. By default the timeout is 60 seconds, unless a custom httpx client is used, in which case this default is not enforced.
+
+ follow_redirects : typing.Optional[bool]
+ Whether the default httpx client follows redirects or not, this is irrelevant if a custom httpx client is passed in.
+
+ httpx_client : typing.Optional[httpx.Client]
+ The httpx client to use for making requests, a preconfigured client is used by default, however this is useful should you want to pass in any custom httpx configuration.
+
+ Examples
+ --------
+ from skyflow import SkyflowAuth
+
+ client = SkyflowAuth(
+ token="YOUR_TOKEN",
+ )
+ """
+
+ def __init__(
+ self,
+ *,
+ base_url: typing.Optional[str] = None,
+ environment: SkyflowAuthEnvironment = SkyflowAuthEnvironment.PRODUCTION,
+ token: typing.Optional[typing.Union[str, typing.Callable[[], str]]] = None,
+ headers: typing.Optional[typing.Dict[str, str]] = None,
+ timeout: typing.Optional[float] = None,
+ follow_redirects: typing.Optional[bool] = True,
+ httpx_client: typing.Optional[httpx.Client] = None,
+ ):
+ _defaulted_timeout = (
+ timeout if timeout is not None else 60 if httpx_client is None else httpx_client.timeout.read
+ )
+ self._client_wrapper = SyncClientWrapper(
+ base_url=_get_base_url(base_url=base_url, environment=environment),
+ token=token,
+ headers=headers,
+ httpx_client=httpx_client
+ if httpx_client is not None
+ else httpx.Client(timeout=_defaulted_timeout, follow_redirects=follow_redirects)
+ if follow_redirects is not None
+ else httpx.Client(timeout=_defaulted_timeout),
+ timeout=_defaulted_timeout,
+ )
+ self.authentication = AuthenticationClient(client_wrapper=self._client_wrapper)
+
+
+class AsyncSkyflowAuth:
+ """
+ Use this class to access the different functions within the SDK. You can instantiate any number of clients with different configuration that will propagate to these functions.
+
+ Parameters
+ ----------
+ base_url : typing.Optional[str]
+ The base url to use for requests from the client.
+
+ environment : SkyflowAuthEnvironment
+ The environment to use for requests from the client. from .environment import SkyflowAuthEnvironment
+
+
+
+ Defaults to SkyflowAuthEnvironment.PRODUCTION
+
+
+
+ token : typing.Optional[typing.Union[str, typing.Callable[[], str]]]
+ headers : typing.Optional[typing.Dict[str, str]]
+ Additional headers to send with every request.
+
+ timeout : typing.Optional[float]
+ The timeout to be used, in seconds, for requests. By default the timeout is 60 seconds, unless a custom httpx client is used, in which case this default is not enforced.
+
+ follow_redirects : typing.Optional[bool]
+ Whether the default httpx client follows redirects or not, this is irrelevant if a custom httpx client is passed in.
+
+ httpx_client : typing.Optional[httpx.AsyncClient]
+ The httpx client to use for making requests, a preconfigured client is used by default, however this is useful should you want to pass in any custom httpx configuration.
+
+ Examples
+ --------
+ from skyflow import AsyncSkyflowAuth
+
+ client = AsyncSkyflowAuth(
+ token="YOUR_TOKEN",
+ )
+ """
+
+ def __init__(
+ self,
+ *,
+ base_url: typing.Optional[str] = None,
+ environment: SkyflowAuthEnvironment = SkyflowAuthEnvironment.PRODUCTION,
+ token: typing.Optional[typing.Union[str, typing.Callable[[], str]]] = None,
+ headers: typing.Optional[typing.Dict[str, str]] = None,
+ timeout: typing.Optional[float] = None,
+ follow_redirects: typing.Optional[bool] = True,
+ httpx_client: typing.Optional[httpx.AsyncClient] = None,
+ ):
+ _defaulted_timeout = (
+ timeout if timeout is not None else 60 if httpx_client is None else httpx_client.timeout.read
+ )
+ self._client_wrapper = AsyncClientWrapper(
+ base_url=_get_base_url(base_url=base_url, environment=environment),
+ token=token,
+ headers=headers,
+ httpx_client=httpx_client
+ if httpx_client is not None
+ else httpx.AsyncClient(timeout=_defaulted_timeout, follow_redirects=follow_redirects)
+ if follow_redirects is not None
+ else httpx.AsyncClient(timeout=_defaulted_timeout),
+ timeout=_defaulted_timeout,
+ )
+ self.authentication = AsyncAuthenticationClient(client_wrapper=self._client_wrapper)
+
+
+def _get_base_url(*, base_url: typing.Optional[str] = None, environment: SkyflowAuthEnvironment) -> str:
+ if base_url is not None:
+ return base_url
+ elif environment is not None:
+ return environment.value
+ else:
+ raise Exception("Please pass in either base_url or environment to construct the client")
diff --git a/skyflow/generated/rest/core/__init__.py b/common/generated/rest/core/__init__.py
similarity index 100%
rename from skyflow/generated/rest/core/__init__.py
rename to common/generated/rest/core/__init__.py
diff --git a/skyflow/generated/rest/core/api_error.py b/common/generated/rest/core/api_error.py
similarity index 100%
rename from skyflow/generated/rest/core/api_error.py
rename to common/generated/rest/core/api_error.py
diff --git a/common/generated/rest/core/client_wrapper.py b/common/generated/rest/core/client_wrapper.py
new file mode 100644
index 00000000..9588d762
--- /dev/null
+++ b/common/generated/rest/core/client_wrapper.py
@@ -0,0 +1,86 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import httpx
+from .http_client import AsyncHttpClient, HttpClient
+
+
+class BaseClientWrapper:
+ def __init__(
+ self,
+ *,
+ token: typing.Optional[typing.Union[str, typing.Callable[[], str]]] = None,
+ headers: typing.Optional[typing.Dict[str, str]] = None,
+ base_url: str,
+ timeout: typing.Optional[float] = None,
+ ):
+ self._token = token
+ self._headers = headers
+ self._base_url = base_url
+ self._timeout = timeout
+
+ def get_headers(self) -> typing.Dict[str, str]:
+ headers: typing.Dict[str, str] = {
+ "X-Fern-Language": "Python",
+ "X-Fern-SDK-Name": "skyflow.generated.rest",
+ "X-Fern-SDK-Version": "0.0.9",
+ **(self.get_custom_headers() or {}),
+ }
+ token = self._get_token()
+ if token is not None:
+ headers["Authorization"] = f"Bearer {token}"
+ return headers
+
+ def _get_token(self) -> typing.Optional[str]:
+ if isinstance(self._token, str) or self._token is None:
+ return self._token
+ else:
+ return self._token()
+
+ def get_custom_headers(self) -> typing.Optional[typing.Dict[str, str]]:
+ return self._headers
+
+ def get_base_url(self) -> str:
+ return self._base_url
+
+ def get_timeout(self) -> typing.Optional[float]:
+ return self._timeout
+
+
+class SyncClientWrapper(BaseClientWrapper):
+ def __init__(
+ self,
+ *,
+ token: typing.Optional[typing.Union[str, typing.Callable[[], str]]] = None,
+ headers: typing.Optional[typing.Dict[str, str]] = None,
+ base_url: str,
+ timeout: typing.Optional[float] = None,
+ httpx_client: httpx.Client,
+ ):
+ super().__init__(token=token, headers=headers, base_url=base_url, timeout=timeout)
+ self.httpx_client = HttpClient(
+ httpx_client=httpx_client,
+ base_headers=self.get_headers,
+ base_timeout=self.get_timeout,
+ base_url=self.get_base_url,
+ )
+
+
+class AsyncClientWrapper(BaseClientWrapper):
+ def __init__(
+ self,
+ *,
+ token: typing.Optional[typing.Union[str, typing.Callable[[], str]]] = None,
+ headers: typing.Optional[typing.Dict[str, str]] = None,
+ base_url: str,
+ timeout: typing.Optional[float] = None,
+ httpx_client: httpx.AsyncClient,
+ ):
+ super().__init__(token=token, headers=headers, base_url=base_url, timeout=timeout)
+ self.httpx_client = AsyncHttpClient(
+ httpx_client=httpx_client,
+ base_headers=self.get_headers,
+ base_timeout=self.get_timeout,
+ base_url=self.get_base_url,
+ )
diff --git a/skyflow/generated/rest/core/datetime_utils.py b/common/generated/rest/core/datetime_utils.py
similarity index 100%
rename from skyflow/generated/rest/core/datetime_utils.py
rename to common/generated/rest/core/datetime_utils.py
diff --git a/skyflow/generated/rest/core/file.py b/common/generated/rest/core/file.py
similarity index 100%
rename from skyflow/generated/rest/core/file.py
rename to common/generated/rest/core/file.py
diff --git a/skyflow/generated/rest/core/force_multipart.py b/common/generated/rest/core/force_multipart.py
similarity index 100%
rename from skyflow/generated/rest/core/force_multipart.py
rename to common/generated/rest/core/force_multipart.py
diff --git a/skyflow/generated/rest/core/http_client.py b/common/generated/rest/core/http_client.py
similarity index 100%
rename from skyflow/generated/rest/core/http_client.py
rename to common/generated/rest/core/http_client.py
diff --git a/skyflow/generated/rest/core/http_response.py b/common/generated/rest/core/http_response.py
similarity index 100%
rename from skyflow/generated/rest/core/http_response.py
rename to common/generated/rest/core/http_response.py
diff --git a/skyflow/generated/rest/core/jsonable_encoder.py b/common/generated/rest/core/jsonable_encoder.py
similarity index 100%
rename from skyflow/generated/rest/core/jsonable_encoder.py
rename to common/generated/rest/core/jsonable_encoder.py
diff --git a/skyflow/generated/rest/core/pydantic_utilities.py b/common/generated/rest/core/pydantic_utilities.py
similarity index 100%
rename from skyflow/generated/rest/core/pydantic_utilities.py
rename to common/generated/rest/core/pydantic_utilities.py
diff --git a/skyflow/generated/rest/core/query_encoder.py b/common/generated/rest/core/query_encoder.py
similarity index 100%
rename from skyflow/generated/rest/core/query_encoder.py
rename to common/generated/rest/core/query_encoder.py
diff --git a/skyflow/generated/rest/core/remove_none_from_dict.py b/common/generated/rest/core/remove_none_from_dict.py
similarity index 100%
rename from skyflow/generated/rest/core/remove_none_from_dict.py
rename to common/generated/rest/core/remove_none_from_dict.py
diff --git a/skyflow/generated/rest/core/request_options.py b/common/generated/rest/core/request_options.py
similarity index 100%
rename from skyflow/generated/rest/core/request_options.py
rename to common/generated/rest/core/request_options.py
diff --git a/skyflow/generated/rest/core/serialization.py b/common/generated/rest/core/serialization.py
similarity index 100%
rename from skyflow/generated/rest/core/serialization.py
rename to common/generated/rest/core/serialization.py
diff --git a/common/generated/rest/environment.py b/common/generated/rest/environment.py
new file mode 100644
index 00000000..b1c13812
--- /dev/null
+++ b/common/generated/rest/environment.py
@@ -0,0 +1,8 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import enum
+
+
+class SkyflowAuthEnvironment(enum.Enum):
+ PRODUCTION = "https://manage.skyflowapis.com"
+ SANDBOX = "https://manage.skyflowapis-preview.com"
diff --git a/common/generated/rest/errors/__init__.py b/common/generated/rest/errors/__init__.py
new file mode 100644
index 00000000..fdf6196c
--- /dev/null
+++ b/common/generated/rest/errors/__init__.py
@@ -0,0 +1,9 @@
+# This file was auto-generated by Fern from our API Definition.
+
+# isort: skip_file
+
+from .bad_request_error import BadRequestError
+from .not_found_error import NotFoundError
+from .unauthorized_error import UnauthorizedError
+
+__all__ = ["BadRequestError", "NotFoundError", "UnauthorizedError"]
diff --git a/common/generated/rest/errors/bad_request_error.py b/common/generated/rest/errors/bad_request_error.py
new file mode 100644
index 00000000..c5d0db48
--- /dev/null
+++ b/common/generated/rest/errors/bad_request_error.py
@@ -0,0 +1,14 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+from ..core.api_error import ApiError
+
+
+class BadRequestError(ApiError):
+ def __init__(
+ self,
+ body: typing.Dict[str, typing.Optional[typing.Any]],
+ headers: typing.Optional[typing.Dict[str, str]] = None,
+ ):
+ super().__init__(status_code=400, headers=headers, body=body)
diff --git a/common/generated/rest/errors/not_found_error.py b/common/generated/rest/errors/not_found_error.py
new file mode 100644
index 00000000..66307415
--- /dev/null
+++ b/common/generated/rest/errors/not_found_error.py
@@ -0,0 +1,14 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+from ..core.api_error import ApiError
+
+
+class NotFoundError(ApiError):
+ def __init__(
+ self,
+ body: typing.Dict[str, typing.Optional[typing.Any]],
+ headers: typing.Optional[typing.Dict[str, str]] = None,
+ ):
+ super().__init__(status_code=404, headers=headers, body=body)
diff --git a/common/generated/rest/errors/unauthorized_error.py b/common/generated/rest/errors/unauthorized_error.py
new file mode 100644
index 00000000..3d58c2e6
--- /dev/null
+++ b/common/generated/rest/errors/unauthorized_error.py
@@ -0,0 +1,14 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+from ..core.api_error import ApiError
+
+
+class UnauthorizedError(ApiError):
+ def __init__(
+ self,
+ body: typing.Dict[str, typing.Optional[typing.Any]],
+ headers: typing.Optional[typing.Dict[str, str]] = None,
+ ):
+ super().__init__(status_code=401, headers=headers, body=body)
diff --git a/skyflow/generated/rest/py.typed b/common/generated/rest/py.typed
similarity index 100%
rename from skyflow/generated/rest/py.typed
rename to common/generated/rest/py.typed
diff --git a/common/generated/rest/types/__init__.py b/common/generated/rest/types/__init__.py
new file mode 100644
index 00000000..8a9140f5
--- /dev/null
+++ b/common/generated/rest/types/__init__.py
@@ -0,0 +1,9 @@
+# This file was auto-generated by Fern from our API Definition.
+
+# isort: skip_file
+
+from .googlerpc_status import GooglerpcStatus
+from .protobuf_any import ProtobufAny
+from .v_1_get_auth_token_response import V1GetAuthTokenResponse
+
+__all__ = ["GooglerpcStatus", "ProtobufAny", "V1GetAuthTokenResponse"]
diff --git a/skyflow/generated/rest/types/googlerpc_status.py b/common/generated/rest/types/googlerpc_status.py
similarity index 100%
rename from skyflow/generated/rest/types/googlerpc_status.py
rename to common/generated/rest/types/googlerpc_status.py
diff --git a/skyflow/generated/rest/types/protobuf_any.py b/common/generated/rest/types/protobuf_any.py
similarity index 100%
rename from skyflow/generated/rest/types/protobuf_any.py
rename to common/generated/rest/types/protobuf_any.py
diff --git a/skyflow/generated/rest/types/v_1_get_auth_token_response.py b/common/generated/rest/types/v_1_get_auth_token_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_get_auth_token_response.py
rename to common/generated/rest/types/v_1_get_auth_token_response.py
diff --git a/common/generated/rest/version.py b/common/generated/rest/version.py
new file mode 100644
index 00000000..4f3bb47a
--- /dev/null
+++ b/common/generated/rest/version.py
@@ -0,0 +1,6 @@
+# NOTE: hand-patched, not Fern-generated content. Fern originally emitted a runtime
+# metadata.version("skyflow.generated.rest") lookup here, but this code is bundled into the
+# skyflow (v2) and v3 wheels via a build_py hook rather than published under that distribution
+# name, so the lookup always raised PackageNotFoundError on import. Hardcoded until the Fern
+# generator config (skyflow-fern-config) is updated to stop emitting a runtime lookup here.
+__version__ = "0.0.9"
diff --git a/common/service_account/__init__.py b/common/service_account/__init__.py
new file mode 100644
index 00000000..58a05b45
--- /dev/null
+++ b/common/service_account/__init__.py
@@ -0,0 +1 @@
+from ._utils import generate_bearer_token, generate_bearer_token_from_creds, is_expired, generate_signed_data_tokens, generate_signed_data_tokens_from_creds
diff --git a/common/service_account/_utils.py b/common/service_account/_utils.py
new file mode 100644
index 00000000..1c9bf945
--- /dev/null
+++ b/common/service_account/_utils.py
@@ -0,0 +1,247 @@
+import json
+import datetime
+import re
+import time
+import jwt
+from urllib.parse import urlparse
+from common.errors import SkyflowError
+from common.service_account.client.auth_client import AuthClient
+from common.utils.logger import log_info, log_error_log
+from common.utils import get_base_url, format_scope, SkyflowMessages
+from common.utils.constants import JWT, CredentialField, JwtField, OptionField, ResponseField
+from common.generated.rest.errors.unauthorized_error import UnauthorizedError
+from common.utils import is_valid_url
+from common.utils.constants import CTX_KEY_REGEX
+
+
+invalid_input_error_code = SkyflowMessages.ErrorCodes.INVALID_INPUT.value
+
+_CTX_KEY_PATTERN = re.compile(CTX_KEY_REGEX)
+
+_SNAKE_TO_CAMEL_CRED_MAP = {
+ 'private_key': CredentialField.PRIVATE_KEY,
+ 'client_id': CredentialField.CLIENT_ID,
+ 'key_id': CredentialField.KEY_ID,
+ 'token_uri': CredentialField.TOKEN_URI,
+ 'client_name': CredentialField.CLIENT_NAME,
+}
+
+
+def _normalize_credentials(credentials):
+ return {_SNAKE_TO_CAMEL_CRED_MAP.get(k, k): v for k, v in credentials.items()}
+
+
+def _validate_and_resolve_ctx(ctx):
+ """Validate ctx value and return resolved value for JWT claims.
+ Returns None if ctx should be omitted, the value if valid, or raises SkyflowError if invalid.
+ """
+ if ctx is None:
+ return None
+ if isinstance(ctx, str):
+ if ctx.strip() == '':
+ return None
+ return ctx
+ if isinstance(ctx, dict):
+ if len(ctx) == 0:
+ return None
+ for key in ctx:
+ if not isinstance(key, str) or not _CTX_KEY_PATTERN.match(key):
+ raise SkyflowError(
+ SkyflowMessages.Error.INVALID_CTX_MAP_KEY.value.format(key),
+ invalid_input_error_code
+ )
+ return ctx
+ if isinstance(ctx, (bool, int, float)):
+ return ctx
+ raise SkyflowError(
+ SkyflowMessages.Error.INVALID_CTX_TYPE.value,
+ invalid_input_error_code
+ )
+
+def is_expired(token, logger = None):
+ if token is None:
+ return True
+ if len(token) == 0:
+ log_error_log(SkyflowMessages.ErrorLogs.INVALID_BEARER_TOKEN.value)
+ return True
+
+ try:
+ decoded = jwt.decode(
+ token, options={OptionField.VERIFY_SIGNATURE: False, OptionField.VERIFY_AUD: False})
+ if time.time() >= decoded[JwtField.EXP]:
+ log_info(SkyflowMessages.Info.BEARER_TOKEN_EXPIRED.value, logger)
+ log_error_log(SkyflowMessages.ErrorLogs.INVALID_BEARER_TOKEN.value)
+ return True
+ return False
+ except jwt.ExpiredSignatureError:
+ return True
+ except Exception:
+ log_error_log(SkyflowMessages.Error.JWT_DECODE_ERROR.value, logger)
+ return True
+
+def generate_bearer_token(credentials_file_path, options = None, logger = None):
+ log_info(SkyflowMessages.Info.GET_BEARER_TOKEN_TRIGGERED.value, logger)
+ try:
+ with open(credentials_file_path, 'r') as credentials_file:
+ try:
+ credentials = json.load(credentials_file)
+ except Exception:
+ log_error_log(SkyflowMessages.ErrorLogs.INVALID_CREDENTIALS_FILE.value, logger=logger)
+ raise SkyflowError(SkyflowMessages.Error.FILE_INVALID_JSON.value.format(credentials_file_path), invalid_input_error_code)
+ except SkyflowError:
+ raise
+ except Exception:
+ raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIAL_FILE_PATH.value, invalid_input_error_code)
+ result = get_service_account_token(credentials, options, logger)
+ return result
+
+def generate_bearer_token_from_creds(credentials, options = None, logger = None):
+ log_info(SkyflowMessages.Info.GET_BEARER_TOKEN_TRIGGERED.value, logger)
+ credentials = credentials.strip()
+ try:
+ json_credentials = json.loads(credentials.replace('\n', '\\n'))
+ except Exception:
+ raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIALS_STRING.value, invalid_input_error_code)
+ result = get_service_account_token(json_credentials, options, logger)
+ return result
+
+def get_service_account_token(credentials, options, logger):
+ credentials = _normalize_credentials(credentials)
+ try:
+ private_key = credentials[CredentialField.PRIVATE_KEY]
+ except KeyError:
+ log_error_log(SkyflowMessages.ErrorLogs.PRIVATE_KEY_IS_REQUIRED.value, logger=logger)
+ raise SkyflowError(SkyflowMessages.Error.MISSING_PRIVATE_KEY.value, invalid_input_error_code)
+ try:
+ client_id = credentials[CredentialField.CLIENT_ID]
+ except KeyError:
+ log_error_log(SkyflowMessages.ErrorLogs.CLIENT_ID_IS_REQUIRED.value, logger=logger)
+ raise SkyflowError(SkyflowMessages.Error.MISSING_CLIENT_ID.value, invalid_input_error_code)
+ try:
+ key_id = credentials[CredentialField.KEY_ID]
+ except KeyError:
+ log_error_log(SkyflowMessages.ErrorLogs.KEY_ID_IS_REQUIRED.value, logger=logger)
+ raise SkyflowError(SkyflowMessages.Error.MISSING_KEY_ID.value, invalid_input_error_code)
+ try:
+ token_uri = credentials[CredentialField.TOKEN_URI]
+ except KeyError:
+ log_error_log(SkyflowMessages.ErrorLogs.TOKEN_URI_IS_REQUIRED.value, logger=logger)
+ raise SkyflowError(SkyflowMessages.Error.MISSING_TOKEN_URI.value, invalid_input_error_code)
+
+ if not isinstance(token_uri, str) or not is_valid_url(token_uri):
+ log_error_log(SkyflowMessages.ErrorLogs.INVALID_TOKEN_URI.value, logger=logger)
+ raise SkyflowError(SkyflowMessages.Error.INVALID_TOKEN_URI.value, invalid_input_error_code)
+
+ if options and CredentialField.TOKEN_URI_OPTION in options:
+ token_uri = options[CredentialField.TOKEN_URI_OPTION]
+ if not isinstance(token_uri, str) or not is_valid_url(token_uri):
+ log_error_log(SkyflowMessages.ErrorLogs.INVALID_TOKEN_URI.value, logger=logger)
+ raise SkyflowError(SkyflowMessages.Error.INVALID_TOKEN_URI.value, invalid_input_error_code)
+
+ signed_token = get_signed_jwt(options, client_id, key_id, token_uri, private_key, logger)
+ base_url = get_base_url(token_uri)
+ auth_client = AuthClient(base_url)
+ auth_api = auth_client.get_auth_api()
+
+ formatted_scope = None
+ if options and OptionField.ROLE_IDS in options:
+ formatted_scope = format_scope(options.get(OptionField.ROLE_IDS))
+
+ try:
+ response = auth_api.authentication_service_get_auth_token(assertion = signed_token,
+ grant_type=JWT.GRANT_TYPE_JWT_BEARER,
+ scope=formatted_scope)
+ log_info(SkyflowMessages.Info.GET_BEARER_TOKEN_SUCCESS.value, logger)
+ except UnauthorizedError:
+ log_error_log(SkyflowMessages.ErrorLogs.UNAUTHORIZED_ERROR_IN_GETTING_BEARER_TOKEN.value, logger=logger)
+ raise SkyflowError(SkyflowMessages.Error.UNAUTHORIZED_ERROR_IN_GETTING_BEARER_TOKEN.value, invalid_input_error_code)
+ except Exception:
+ log_error_log(SkyflowMessages.ErrorLogs.FAILED_TO_GET_BEARER_TOKEN.value, logger=logger)
+ raise SkyflowError(SkyflowMessages.Error.FAILED_TO_GET_BEARER_TOKEN.value, invalid_input_error_code)
+ return response.access_token, response.token_type
+
+def get_signed_jwt(options, client_id, key_id, token_uri, private_key, logger):
+ payload = {
+ JwtField.ISS: client_id,
+ JwtField.KEY: key_id,
+ JwtField.AUD: token_uri,
+ JwtField.SUB: client_id,
+ JwtField.EXP: datetime.datetime.utcnow() + datetime.timedelta(minutes=60)
+ }
+ if options and OptionField.CTX in options:
+ resolved_ctx = _validate_and_resolve_ctx(options.get(OptionField.CTX))
+ if resolved_ctx is not None:
+ payload[JwtField.CTX] = resolved_ctx
+ try:
+ return jwt.encode(payload=payload, key=private_key, algorithm=JWT.ALGORITHM_RS256)
+ except Exception:
+ raise SkyflowError(SkyflowMessages.Error.JWT_INVALID_FORMAT.value, invalid_input_error_code)
+
+
+
+def get_signed_tokens(credentials_obj, options):
+ options = options if options is not None else {}
+ credentials_obj = _normalize_credentials(credentials_obj)
+ expiry_time = int(time.time()) + options.get(OptionField.TIME_TO_LIVE, 60)
+ prefix = JWT.SIGNED_TOKEN_PREFIX
+
+ token_uri = credentials_obj.get(CredentialField.TOKEN_URI)
+ if not isinstance(token_uri, str) or not is_valid_url(token_uri):
+ log_error_log(SkyflowMessages.ErrorLogs.INVALID_TOKEN_URI.value)
+ raise SkyflowError(SkyflowMessages.Error.INVALID_TOKEN_URI.value, invalid_input_error_code)
+
+ resolved_ctx = None
+ if OptionField.CTX in options:
+ resolved_ctx = _validate_and_resolve_ctx(options[OptionField.CTX])
+
+ results = []
+ if options and options.get(OptionField.DATA_TOKENS):
+ for token in options[OptionField.DATA_TOKENS]:
+ claims = {
+ JwtField.ISS: JWT.ISSUER_SDK,
+ JwtField.KEY: credentials_obj.get(CredentialField.KEY_ID),
+ JwtField.EXP: expiry_time,
+ JwtField.SUB: credentials_obj.get(CredentialField.CLIENT_ID),
+ JwtField.TOK: token,
+ JwtField.IAT: int(time.time()),
+ }
+ if resolved_ctx is not None:
+ claims[JwtField.CTX] = resolved_ctx
+ private_key = credentials_obj.get(CredentialField.PRIVATE_KEY)
+ try:
+ signed_jwt = jwt.encode(claims, private_key, algorithm=JWT.ALGORITHM_RS256)
+ except Exception:
+ raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIALS.value, invalid_input_error_code)
+ results.append(get_signed_data_token_response_object(prefix + signed_jwt, token))
+ log_info(SkyflowMessages.Info.GET_SIGNED_DATA_TOKEN_SUCCESS.value)
+ return results
+
+
+def generate_signed_data_tokens(credentials_file_path, options):
+ log_info(SkyflowMessages.Info.GET_SIGNED_DATA_TOKENS_TRIGGERED.value)
+ try:
+ with open(credentials_file_path, 'r') as credentials_file:
+ try:
+ credentials = json.load(credentials_file)
+ except Exception:
+ raise SkyflowError(SkyflowMessages.Error.FILE_INVALID_JSON.value.format(credentials_file_path),
+ invalid_input_error_code)
+ except SkyflowError:
+ raise
+ except Exception:
+ raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIAL_FILE_PATH.value, invalid_input_error_code)
+ return get_signed_tokens(credentials, options)
+
+def generate_signed_data_tokens_from_creds(credentials, options):
+ log_info(SkyflowMessages.Info.GET_SIGNED_DATA_TOKENS_TRIGGERED.value)
+ credentials = credentials.strip()
+ try:
+ json_credentials = json.loads(credentials.replace('\n', '\\n'))
+ except Exception:
+ log_error_log(SkyflowMessages.ErrorLogs.INVALID_CREDENTIALS_FILE.value)
+ raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIALS_STRING.value, invalid_input_error_code)
+ return get_signed_tokens(json_credentials, options)
+
+
+def get_signed_data_token_response_object(signed_token, actual_token):
+ return actual_token, signed_token
diff --git a/skyflow/vault/__init__.py b/common/service_account/client/__init__.py
similarity index 100%
rename from skyflow/vault/__init__.py
rename to common/service_account/client/__init__.py
diff --git a/common/service_account/client/auth_client.py b/common/service_account/client/auth_client.py
new file mode 100644
index 00000000..fe7aabe8
--- /dev/null
+++ b/common/service_account/client/auth_client.py
@@ -0,0 +1,13 @@
+from common.generated.rest.client import SkyflowAuth
+from common.utils.constants import OPTIONAL_TOKEN
+
+class AuthClient:
+ def __init__(self, url):
+ self.__url = url
+ self.__api_client = self.initialize_api_client()
+
+ def initialize_api_client(self):
+ return SkyflowAuth(base_url=self.__url, token=OPTIONAL_TOKEN)
+
+ def get_auth_api(self):
+ return self.__api_client.authentication
diff --git a/common/setup.py b/common/setup.py
new file mode 100644
index 00000000..fe37d04d
--- /dev/null
+++ b/common/setup.py
@@ -0,0 +1,27 @@
+'''
+ Copyright (c) 2022 Skyflow, Inc.
+'''
+from setuptools import setup, find_packages
+
+setup(
+ # Never published to PyPI -- bundled into the v2/v3 wheels at build time (see their
+ # setup.py CustomBuildPy). Exists for local dev (`pip install -e ./common`) and tests/contract/.
+ name='skyflow-common',
+ version='0.0.0',
+ author='Skyflow',
+ author_email='service-ops@skyflow.com',
+ description='Internal shared code for the Skyflow Python SDK v2/v3 build variants. Not published independently.',
+ packages=find_packages(where='.', exclude=['tests*']),
+ package_data={
+ 'common.generated.rest': ['py.typed'],
+ },
+ install_requires=[
+ 'pydantic >= 2.0.0',
+ 'typing-extensions >= 4.0.0',
+ 'PyJWT >= 2.12, < 3',
+ 'cryptography >= 44.0.2',
+ 'httpx >= 0.21.2',
+ 'python-dotenv >= 1.1.0, < 2',
+ ],
+ python_requires=">=3.9",
+)
diff --git a/skyflow/vault/client/__init__.py b/common/tests/__init__.py
similarity index 100%
rename from skyflow/vault/client/__init__.py
rename to common/tests/__init__.py
diff --git a/tests/client/__init__.py b/common/tests/client/__init__.py
similarity index 100%
rename from tests/client/__init__.py
rename to common/tests/client/__init__.py
diff --git a/common/tests/client/test_base_skyflow.py b/common/tests/client/test_base_skyflow.py
new file mode 100644
index 00000000..f28e9416
--- /dev/null
+++ b/common/tests/client/test_base_skyflow.py
@@ -0,0 +1,215 @@
+import unittest
+
+from common.errors import SkyflowError
+from common.utils import LogLevel, SkyflowMessages
+from common.utils.logger import Logger
+from common.client.base_skyflow import BaseSkyflow, BaseSkyflowImpl
+from common.client.utils import make_skyflow_class
+
+
+class FakeVaultClient:
+ def __init__(self, config):
+ self._config = dict(config)
+ self.credentials = None
+ self.logger = None
+
+ def get_config(self):
+ return self._config
+
+ def update_config(self, config):
+ self._config.update(config)
+
+ def set_logger(self, log_level, logger):
+ self.logger = logger
+
+ def set_common_skyflow_credentials(self, credentials):
+ self.credentials = credentials
+
+
+class FakeVaultController:
+ def __init__(self, vault_client):
+ self.vault_client = vault_client
+
+
+class FakeConnection:
+ def __init__(self, vault_client):
+ self.vault_client = vault_client
+
+
+class FakeDetect:
+ def __init__(self, vault_client):
+ self.vault_client = vault_client
+
+
+def _noop_validate(logger, config):
+ return True
+
+
+def make_fake_skyflow(with_connections=False, with_detect=False):
+ kwargs = dict(
+ vault_client_cls=FakeVaultClient,
+ vault_controller_cls=FakeVaultController,
+ logger_cls=Logger,
+ default_log_level=LogLevel.ERROR,
+ skyflow_messages=SkyflowMessages,
+ validate_vault_config=_noop_validate,
+ validate_update_vault_config=_noop_validate,
+ validate_log_level=_noop_validate,
+ validate_credentials=_noop_validate,
+ )
+ if with_connections:
+ kwargs.update(
+ connection_cls=FakeConnection,
+ validate_connection_config=_noop_validate,
+ validate_update_connection_config=_noop_validate,
+ )
+ if with_detect:
+ kwargs['detect_cls'] = FakeDetect
+ return make_skyflow_class(**kwargs)
+
+
+VAULT_CONFIG = {"vault_id": "v1", "cluster_id": "c1", "credentials": {"token": "t"}}
+
+
+class TestMakeSkyflowClassBasics(unittest.TestCase):
+ def test_build_returns_instance_of_the_produced_class_not_the_template(self):
+ """Regression pin: build()/builder() must resolve to the specific class produced by
+ make_skyflow_class(), not the shared template -- two variants must never collide."""
+ SkyflowA = make_fake_skyflow()
+ SkyflowB = make_fake_skyflow()
+ client_a = SkyflowA.builder().add_vault_config(VAULT_CONFIG).build()
+ self.assertIsInstance(client_a, SkyflowA)
+ self.assertNotIsInstance(client_a, SkyflowB)
+
+ def test_two_produced_classes_do_not_share_hooks(self):
+ SkyflowWithDetect = make_fake_skyflow(with_detect=True)
+ SkyflowWithoutDetect = make_fake_skyflow(with_detect=False)
+ self.assertIsNotNone(SkyflowWithDetect.Builder._detect_cls)
+ self.assertIsNone(SkyflowWithoutDetect.Builder._detect_cls)
+
+ def test_vault_config_crud(self):
+ Skyflow = make_fake_skyflow()
+ builder = Skyflow.builder()
+ builder.add_vault_config(VAULT_CONFIG)
+ client = builder.build()
+
+ vault_config = client.get_vault_config("v1")
+ self.assertEqual(vault_config.get("vault_id"), "v1")
+
+ updated = dict(VAULT_CONFIG)
+ updated["cluster_id"] = "c2"
+ client.update_vault_config(updated)
+ self.assertEqual(client.get_vault_config("v1").get("cluster_id"), "c2")
+
+ client.remove_vault_config("v1")
+ with self.assertRaises(SkyflowError):
+ client.get_vault_config("v1")
+
+ def test_vault_returns_the_controller(self):
+ Skyflow = make_fake_skyflow()
+ client = Skyflow.builder().add_vault_config(VAULT_CONFIG).build()
+ self.assertIsInstance(client.vault("v1"), FakeVaultController)
+
+ def test_add_skyflow_credentials_and_update(self):
+ Skyflow = make_fake_skyflow()
+ client = Skyflow.builder().add_vault_config(VAULT_CONFIG).build()
+ client.add_skyflow_credentials({"token": "a"})
+ client.update_skyflow_credentials({"token": "b"})
+ # no assertion error means both delegate correctly to the same underlying builder path
+
+ def test_set_get_and_deprecated_update_log_level(self):
+ Skyflow = make_fake_skyflow()
+ client = Skyflow.builder().add_vault_config(VAULT_CONFIG).build()
+ client.set_log_level(LogLevel.INFO)
+ self.assertEqual(client.get_log_level(), LogLevel.INFO)
+
+ client.update_log_level(LogLevel.WARN)
+ self.assertEqual(client.get_log_level(), LogLevel.WARN)
+
+
+class TestConnectionAndDetectGating(unittest.TestCase):
+ def test_connection_methods_do_not_exist_when_connection_cls_not_supplied(self):
+ """connection()/add_connection_config()/etc. come from ConnectionMixin, which
+ make_skyflow_class() only adds to the produced class's bases when connection_cls is
+ given -- when it's not, the methods are genuinely absent, not just guarded."""
+ Skyflow = make_fake_skyflow()
+ client = Skyflow.builder().add_vault_config(VAULT_CONFIG).build()
+ self.assertFalse(hasattr(client, "connection"))
+ with self.assertRaises(AttributeError):
+ client.connection()
+ with self.assertRaises(AttributeError):
+ client.add_connection_config({})
+ with self.assertRaises(AttributeError):
+ client.remove_connection_config("x")
+ with self.assertRaises(AttributeError):
+ client.update_connection_config({})
+ with self.assertRaises(AttributeError):
+ client.get_connection_config("x")
+
+ def test_detect_does_not_exist_when_detect_cls_not_supplied(self):
+ Skyflow = make_fake_skyflow()
+ client = Skyflow.builder().add_vault_config(VAULT_CONFIG).build()
+ self.assertFalse(hasattr(client, "detect"))
+ with self.assertRaises(AttributeError):
+ client.detect()
+
+ def test_connection_config_crud_when_supplied(self):
+ Skyflow = make_fake_skyflow(with_connections=True)
+ connection_config = {"connection_id": "conn1", "connection_url": "https://x", "credentials": {"token": "t"}}
+ client = Skyflow.builder().add_connection_config(connection_config).build()
+
+ self.assertIsInstance(client.connection("conn1"), FakeConnection)
+
+ updated = dict(connection_config)
+ updated["connection_url"] = "https://y"
+ client.update_connection_config(updated)
+ self.assertEqual(client.get_connection_config("conn1").get("connection_url"), "https://y")
+
+ client.remove_connection_config("conn1")
+ with self.assertRaises(SkyflowError):
+ client.get_connection_config("conn1")
+
+ def test_detect_returns_detect_controller_when_supplied(self):
+ Skyflow = make_fake_skyflow(with_detect=True)
+ client = Skyflow.builder().add_vault_config(VAULT_CONFIG).build()
+ self.assertIsInstance(client.detect("v1"), FakeDetect)
+
+ def test_make_skyflow_class_requires_connection_validators_when_connection_cls_given(self):
+ with self.assertRaises(ValueError):
+ make_skyflow_class(
+ vault_client_cls=FakeVaultClient,
+ vault_controller_cls=FakeVaultController,
+ logger_cls=Logger,
+ default_log_level=LogLevel.ERROR,
+ skyflow_messages=SkyflowMessages,
+ validate_vault_config=_noop_validate,
+ validate_update_vault_config=_noop_validate,
+ validate_log_level=_noop_validate,
+ validate_credentials=_noop_validate,
+ connection_cls=FakeConnection,
+ # validate_connection_config/validate_update_connection_config omitted on purpose
+ )
+
+
+class TestBaseSkyflowInterface(unittest.TestCase):
+ def test_using_the_template_builder_directly_raises_with_a_clear_message(self):
+ with self.assertRaises(NotImplementedError) as ctx:
+ BaseSkyflowImpl.Builder()
+ self.assertIn("make_skyflow_class()", str(ctx.exception))
+ self.assertIn("_vault_client_cls", str(ctx.exception))
+
+ def test_make_skyflow_class_produced_builder_constructs_fine(self):
+ Skyflow = make_fake_skyflow()
+ self.assertIsInstance(Skyflow.builder(), BaseSkyflowImpl.Builder)
+
+ def test_instantiating_base_skyflow_impl_directly_raises(self):
+ with self.assertRaises(SkyflowError):
+ BaseSkyflowImpl(None)
+
+ def test_instantiating_the_pure_interface_raises_type_error(self):
+ with self.assertRaises(TypeError):
+ BaseSkyflow()
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/service_account/__init__.py b/common/tests/utils/__init__.py
similarity index 100%
rename from tests/service_account/__init__.py
rename to common/tests/utils/__init__.py
diff --git a/tests/utils/__init__.py b/common/tests/utils/validations/__init__.py
similarity index 100%
rename from tests/utils/__init__.py
rename to common/tests/utils/validations/__init__.py
diff --git a/common/tests/utils/validations/test__validations.py b/common/tests/utils/validations/test__validations.py
new file mode 100644
index 00000000..b10fc008
--- /dev/null
+++ b/common/tests/utils/validations/test__validations.py
@@ -0,0 +1,169 @@
+import unittest
+
+from common.errors import SkyflowError
+from common.utils import SkyflowMessages, LogLevel, Env
+from common.utils.validations import (
+ validate_vault_config,
+ validate_update_vault_config,
+ validate_credentials,
+ validate_log_level,
+)
+
+VALID_VAULT_CONFIG = {
+ "vault_id": "vault123",
+ "cluster_id": "cluster1",
+ "env": Env.PROD,
+ "credentials": {"api_key": "sky-abcde-" + "f" * 32},
+}
+
+
+class FakeMessages:
+ """Stand-in message catalog to confirm validate_vault_config/etc. actually use the
+ `messages` param passed in, rather than silently falling back to common's own."""
+
+ class Error:
+ class _M:
+ def __init__(self, text):
+ self._text = text
+
+ @property
+ def value(self):
+ return self._text
+
+ def format(self, *args, **kwargs):
+ return self._text
+
+ EMPTY_VAULT_ID = _M("FAKE: empty vault id")
+ INVALID_VAULT_ID = _M("FAKE: invalid vault id")
+ EMPTY_CLUSTER_ID = _M("FAKE: empty cluster id")
+ INVALID_CLUSTER_ID = _M("FAKE: invalid cluster id")
+ EMPTY_CREDENTIALS = _M("FAKE: empty credentials")
+ INVALID_ENV = _M("FAKE: invalid env")
+ INVALID_KEY = _M("FAKE: invalid key")
+ INVALID_LOG_LEVEL = _M("FAKE: invalid log level")
+ INVALID_CREDENTIALS = _M("FAKE: invalid credentials")
+ INVALID_CREDENTIALS_IN_CONFIG = _M("FAKE: invalid credentials in config")
+
+ class ErrorLogs:
+ class _M:
+ def __init__(self, text):
+ self._text = text
+
+ @property
+ def value(self):
+ return self._text
+
+ VAULTID_IS_REQUIRED = _M("fake log")
+ CLUSTER_ID_IS_REQUIRED = _M("fake log")
+ CONNECTION_ID_IS_REQUIRED = _M("fake log")
+ INVALID_CONNECTION_URL = _M("fake log")
+ EMPTY_VAULTID = _M("fake log")
+ EMPTY_CLUSTER_ID = _M("fake log")
+ EMPTY_CONNECTION_ID = _M("fake log")
+ EMPTY_CONNECTION_URL = _M("fake log")
+ EMPTY_CREDENTIALS_PATH = _M("fake log")
+ EMPTY_CREDENTIALS_STRING = _M("fake log")
+ EMPTY_TOKEN_VALUE = _M("fake log")
+ EMPTY_API_KEY_VALUE = _M("fake log")
+ INVALID_KEY = _M("fake log")
+ INVALID_LOG_LEVEL = _M("fake log")
+ ENV_IS_REQUIRED = _M("fake log")
+
+ class Info:
+ class _M:
+ def __init__(self, text):
+ self._text = text
+
+ @property
+ def value(self):
+ return self._text
+
+ VALIDATING_VAULT_CONFIG = _M("fake info")
+
+
+class TestValidateVaultConfig(unittest.TestCase):
+ def test_valid_config_passes(self):
+ self.assertTrue(validate_vault_config(None, dict(VALID_VAULT_CONFIG)))
+
+ def test_missing_vault_id_raises(self):
+ config = dict(VALID_VAULT_CONFIG)
+ del config["vault_id"]
+ with self.assertRaises(SkyflowError):
+ validate_vault_config(None, config)
+
+ def test_unknown_key_raises(self):
+ config = dict(VALID_VAULT_CONFIG)
+ config["unexpected"] = True
+ with self.assertRaises(SkyflowError):
+ validate_vault_config(None, config)
+
+ def test_empty_credentials_raises(self):
+ config = dict(VALID_VAULT_CONFIG)
+ config["credentials"] = {}
+ with self.assertRaises(SkyflowError):
+ validate_vault_config(None, config)
+
+ def test_credentials_are_validated(self):
+ config = dict(VALID_VAULT_CONFIG)
+ config["credentials"] = {"api_key": "not-a-valid-key"}
+ with self.assertRaises(SkyflowError):
+ validate_vault_config(None, config)
+
+ def test_uses_injected_messages_for_raised_error(self):
+ config = dict(VALID_VAULT_CONFIG)
+ del config["vault_id"]
+ with self.assertRaises(SkyflowError) as ctx:
+ validate_vault_config(None, config, messages=FakeMessages)
+ self.assertIn("FAKE", ctx.exception.message)
+
+ def test_defaults_to_common_messages_when_not_injected(self):
+ config = dict(VALID_VAULT_CONFIG)
+ del config["vault_id"]
+ with self.assertRaises(SkyflowError) as ctx:
+ validate_vault_config(None, config)
+ self.assertEqual(ctx.exception.message, SkyflowMessages.Error.INVALID_VAULT_ID.value)
+
+
+class TestValidateUpdateVaultConfig(unittest.TestCase):
+ def test_valid_update_passes(self):
+ self.assertTrue(validate_update_vault_config(None, dict(VALID_VAULT_CONFIG)))
+
+ def test_credentials_required_on_update(self):
+ """Unlike validate_vault_config, credentials are mandatory here."""
+ config = dict(VALID_VAULT_CONFIG)
+ del config["credentials"]
+ with self.assertRaises(SkyflowError):
+ validate_update_vault_config(None, config)
+
+ def test_uses_injected_messages(self):
+ config = dict(VALID_VAULT_CONFIG)
+ del config["credentials"]
+ with self.assertRaises(SkyflowError) as ctx:
+ validate_update_vault_config(None, config, messages=FakeMessages)
+ self.assertIn("FAKE", ctx.exception.message)
+
+
+class TestValidateCredentialsMessageInjection(unittest.TestCase):
+ def test_uses_injected_messages(self):
+ with self.assertRaises(SkyflowError) as ctx:
+ validate_credentials(None, {}, messages=FakeMessages)
+ self.assertIn("FAKE", ctx.exception.message)
+
+ def test_defaults_to_common_messages(self):
+ with self.assertRaises(SkyflowError) as ctx:
+ validate_credentials(None, {})
+ self.assertEqual(ctx.exception.message, SkyflowMessages.Error.INVALID_CREDENTIALS.value)
+
+
+class TestValidateLogLevelMessageInjection(unittest.TestCase):
+ def test_valid_log_level_passes(self):
+ validate_log_level(None, LogLevel.INFO) # should not raise
+
+ def test_uses_injected_messages(self):
+ with self.assertRaises(SkyflowError) as ctx:
+ validate_log_level(None, "not-a-log-level", messages=FakeMessages)
+ self.assertIn("FAKE", ctx.exception.message)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/utils/logger/__init__.py b/common/tests/vault/__init__.py
similarity index 100%
rename from tests/utils/logger/__init__.py
rename to common/tests/vault/__init__.py
diff --git a/tests/utils/validations/__init__.py b/common/tests/vault/data/__init__.py
similarity index 100%
rename from tests/utils/validations/__init__.py
rename to common/tests/vault/data/__init__.py
diff --git a/common/tests/vault/data/test_base_insert.py b/common/tests/vault/data/test_base_insert.py
new file mode 100644
index 00000000..f49da1a8
--- /dev/null
+++ b/common/tests/vault/data/test_base_insert.py
@@ -0,0 +1,42 @@
+import unittest
+
+from common.vault.data import BaseInsertRequest, BaseInsertResponse
+
+
+class TestBaseInsertRequest(unittest.TestCase):
+ def test_table_and_values_are_required(self):
+ with self.assertRaises(TypeError):
+ BaseInsertRequest()
+
+ def test_upsert_defaults_to_none(self):
+ request = BaseInsertRequest(table="t1", values=[{"values": {"a": 1}}])
+ self.assertIsNone(request.upsert)
+
+ def test_construction(self):
+ request = BaseInsertRequest(table="t1", values=[{"values": {"a": 1}}], upsert="upsert_val")
+ self.assertEqual(request.table, "t1")
+ self.assertEqual(request.values, [{"values": {"a": 1}}])
+ self.assertEqual(request.upsert, "upsert_val")
+
+
+class TestBaseInsertResponse(unittest.TestCase):
+ def test_defaults(self):
+ response = BaseInsertResponse()
+ self.assertIsNone(response.inserted_fields)
+ self.assertIsNone(response.errors)
+
+ def test_repr_uses_subclass_name(self):
+ class InsertResponse(BaseInsertResponse):
+ pass
+
+ response = InsertResponse(inserted_fields=[{"skyflow_id": "id1"}], errors=None)
+ self.assertIn("InsertResponse", repr(response))
+ self.assertNotIn("BaseInsertResponse", repr(response))
+
+ def test_str_matches_repr(self):
+ response = BaseInsertResponse(inserted_fields=[], errors=[])
+ self.assertEqual(str(response), repr(response))
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/common/tests/vault/test_base_vault_client.py b/common/tests/vault/test_base_vault_client.py
new file mode 100644
index 00000000..b02acbf2
--- /dev/null
+++ b/common/tests/vault/test_base_vault_client.py
@@ -0,0 +1,289 @@
+import unittest
+from unittest.mock import patch, MagicMock
+
+from common.vault.base_vault_client import BaseVaultClient
+
+CONFIG = {
+ "credentials": "some_credentials",
+ "cluster_id": "test_cluster_id",
+ "env": "test_env",
+ "vault_id": "test_vault_id",
+ "roles": ["role_id_1", "role_id_2"],
+ "ctx": "context"
+}
+
+CREDENTIALS_WITH_API_KEY = {"api_key": "dummy_api_key"}
+CREDENTIALS_WITH_TOKEN = {"token": "dummy_static_token"}
+CREDENTIALS_WITH_PATH = {"path": "/some/path/credentials.json"}
+CREDENTIALS_WITH_STRING = {"credentials_string": '{"clientID": "x"}'}
+
+
+class DummyVaultClient(BaseVaultClient):
+ """Minimal concrete subclass -- exercises BaseVaultClient's shared logic without any
+ variant-specific generated-API wiring. Mirrors what test__client.py used to test directly
+ against v2's VaultClient, before initialize_client_configuration/get_bearer_token moved here.
+
+ resolve_vault_url is a real (if trivial) implementation here, not just a stub, because
+ BaseVaultClient.initialize_client_configuration() now delegates URL construction to it --
+ each variant's derivation differs (see BaseVaultClient.resolve_vault_url's docstring), so
+ there's no generic default to fall back on."""
+
+ def resolve_vault_url(self, cluster_id, env, vault_id, logger=None):
+ return "https://test-vault-url.com"
+
+ def initialize_api_client(self, vault_url, bearer_token):
+ self._api_client = MagicMock()
+
+
+class TestBaseVaultClient(unittest.TestCase):
+ def setUp(self):
+ self.vault_client = DummyVaultClient(dict(CONFIG))
+
+ # ------------------------------------------------------------------ #
+ # Basic setters / getters
+ # ------------------------------------------------------------------ #
+
+ def test_set_common_skyflow_credentials(self):
+ credentials = {"api_key": "dummy_api_key"}
+ self.vault_client.set_common_skyflow_credentials(credentials)
+ self.assertEqual(self.vault_client.get_common_skyflow_credentials(), credentials)
+
+ def test_set_logger(self):
+ mock_logger = MagicMock()
+ self.vault_client.set_logger("INFO", mock_logger)
+ self.assertEqual(self.vault_client.get_log_level(), "INFO")
+ self.assertEqual(self.vault_client.get_logger(), mock_logger)
+
+ def test_get_vault_id(self):
+ self.assertEqual(self.vault_client.get_vault_id(), CONFIG["vault_id"])
+
+ def test_get_config(self):
+ self.assertEqual(self.vault_client.get_config(), CONFIG)
+
+ # ------------------------------------------------------------------ #
+ # initialize_client_configuration — first call (slow path)
+ # ------------------------------------------------------------------ #
+
+ @patch("common.vault.base_vault_client.get_credentials")
+ @patch.object(DummyVaultClient, "resolve_vault_url")
+ @patch.object(DummyVaultClient, "initialize_api_client")
+ def test_initialize_client_configuration_first_call(
+ self, mock_init_api_client, mock_resolve_vault_url, mock_get_credentials
+ ):
+ mock_get_credentials.return_value = CREDENTIALS_WITH_API_KEY
+ mock_resolve_vault_url.return_value = "https://test-vault-url.com"
+
+ self.vault_client.initialize_client_configuration()
+
+ mock_get_credentials.assert_called_once_with(
+ CONFIG["credentials"], None, logger=None
+ )
+ mock_resolve_vault_url.assert_called_once_with(
+ CONFIG["cluster_id"], CONFIG["env"], CONFIG["vault_id"], logger=None
+ )
+ mock_init_api_client.assert_called_once()
+
+ # ------------------------------------------------------------------ #
+ # initialize_client_configuration — fast path (static token)
+ # ------------------------------------------------------------------ #
+
+ @patch("common.vault.base_vault_client.get_credentials")
+ @patch.object(DummyVaultClient, "resolve_vault_url")
+ def test_initialize_client_configuration_fast_path_api_key(
+ self, mock_resolve_vault_url, mock_get_credentials
+ ):
+ """Once initialized with api_key, subsequent calls skip all work."""
+ mock_get_credentials.return_value = CREDENTIALS_WITH_API_KEY
+ mock_resolve_vault_url.return_value = "https://test-vault-url.com"
+
+ self.vault_client.initialize_client_configuration() # first call — slow path
+ mock_get_credentials.reset_mock()
+ mock_resolve_vault_url.reset_mock()
+
+ self.vault_client.initialize_client_configuration() # second call — fast path
+
+ mock_get_credentials.assert_not_called()
+ mock_resolve_vault_url.assert_not_called()
+
+ @patch("common.vault.base_vault_client.get_credentials")
+ @patch.object(DummyVaultClient, "resolve_vault_url")
+ def test_initialize_client_configuration_fast_path_static_token(
+ self, mock_resolve_vault_url, mock_get_credentials
+ ):
+ """Once initialized with a static token, subsequent calls skip all work."""
+ mock_get_credentials.return_value = CREDENTIALS_WITH_TOKEN
+ mock_resolve_vault_url.return_value = "https://test-vault-url.com"
+
+ self.vault_client.initialize_client_configuration()
+ mock_get_credentials.reset_mock()
+ mock_resolve_vault_url.reset_mock()
+
+ self.vault_client.initialize_client_configuration()
+
+ mock_get_credentials.assert_not_called()
+ mock_resolve_vault_url.assert_not_called()
+
+ # ------------------------------------------------------------------ #
+ # initialize_client_configuration — fast path (service account)
+ # ------------------------------------------------------------------ #
+
+ @patch("common.vault.base_vault_client.is_expired", return_value=False)
+ @patch("common.vault.base_vault_client.get_credentials")
+ @patch.object(DummyVaultClient, "resolve_vault_url")
+ @patch.object(DummyVaultClient, "initialize_api_client")
+ def test_initialize_client_configuration_fast_path_valid_sa_token(
+ self, mock_init_api_client, mock_resolve_vault_url, mock_get_credentials, mock_is_expired
+ ):
+ """Service account with a still-valid token skips get_bearer_token entirely."""
+ mock_get_credentials.return_value = CREDENTIALS_WITH_PATH
+ mock_resolve_vault_url.return_value = "https://test-vault-url.com"
+
+ # Seed the cached bearer token as if first call already ran
+ self.vault_client._api_client = MagicMock()
+ self.vault_client._is_static_token = False
+ self.vault_client._bearer_token = "cached_sa_token"
+ self.vault_client._credentials = CREDENTIALS_WITH_PATH
+
+ self.vault_client.initialize_client_configuration()
+
+ mock_get_credentials.assert_not_called()
+ mock_resolve_vault_url.assert_not_called()
+ mock_init_api_client.assert_not_called()
+
+ # ------------------------------------------------------------------ #
+ # initialize_client_configuration — token expiry (no client reinit)
+ # ------------------------------------------------------------------ #
+
+ @patch("common.vault.base_vault_client.generate_bearer_token", return_value=("new_sa_token", None))
+ @patch("common.vault.base_vault_client.is_expired", return_value=True)
+ @patch("common.vault.base_vault_client.get_credentials")
+ @patch.object(DummyVaultClient, "resolve_vault_url")
+ @patch.object(DummyVaultClient, "initialize_api_client")
+ def test_initialize_client_configuration_expired_token_no_reinit(
+ self, mock_init_api_client, mock_resolve_vault_url, mock_get_credentials,
+ mock_is_expired, mock_generate_bearer_token
+ ):
+ """Expired service account token is regenerated in-place; the api client is NOT recreated."""
+ mock_get_credentials.return_value = CREDENTIALS_WITH_PATH
+ mock_resolve_vault_url.return_value = "https://test-vault-url.com"
+
+ # Client already initialized — simulate warm state with an expired token
+ self.vault_client._api_client = MagicMock()
+ self.vault_client._is_static_token = False
+ self.vault_client._bearer_token = "expired_sa_token"
+ self.vault_client._credentials = CREDENTIALS_WITH_PATH
+
+ self.vault_client.initialize_client_configuration()
+
+ # Token was regenerated
+ mock_generate_bearer_token.assert_called_once()
+ self.assertEqual(self.vault_client._bearer_token, "new_sa_token")
+ # api client was NOT recreated
+ mock_init_api_client.assert_not_called()
+
+ # ------------------------------------------------------------------ #
+ # initialize_client_configuration — config update forces reinit
+ # ------------------------------------------------------------------ #
+
+ @patch("common.vault.base_vault_client.get_credentials")
+ @patch.object(DummyVaultClient, "resolve_vault_url")
+ @patch.object(DummyVaultClient, "initialize_api_client")
+ def test_initialize_client_configuration_reinit_after_update_config(
+ self, mock_init_api_client, mock_resolve_vault_url, mock_get_credentials
+ ):
+ """update_config() marks the client stale; next call must recreate it."""
+ mock_get_credentials.return_value = CREDENTIALS_WITH_API_KEY
+ mock_resolve_vault_url.return_value = "https://test-vault-url.com"
+
+ # Simulate already-initialized client
+ self.vault_client._api_client = MagicMock()
+ self.vault_client._is_static_token = True
+
+ self.vault_client.update_config({"cluster_id": "new_cluster"})
+ self.vault_client.initialize_client_configuration()
+
+ mock_get_credentials.assert_called_once()
+ mock_resolve_vault_url.assert_called_once()
+ mock_init_api_client.assert_called_once()
+
+ # ------------------------------------------------------------------ #
+ # get_bearer_token
+ # ------------------------------------------------------------------ #
+
+ def test_get_bearer_token_with_api_key(self):
+ result = self.vault_client.get_bearer_token(CREDENTIALS_WITH_API_KEY)
+ self.assertEqual(result, "dummy_api_key")
+
+ def test_get_bearer_token_with_static_token(self):
+ result = self.vault_client.get_bearer_token(CREDENTIALS_WITH_TOKEN)
+ self.assertEqual(result, "dummy_static_token")
+
+ @patch("common.vault.base_vault_client.generate_bearer_token", return_value=("sa_token", None))
+ def test_get_bearer_token_generates_from_path_on_first_call(self, mock_generate):
+ result = self.vault_client.get_bearer_token(CREDENTIALS_WITH_PATH)
+ mock_generate.assert_called_once()
+ self.assertEqual(result, "sa_token")
+ self.assertEqual(self.vault_client._bearer_token, "sa_token")
+
+ @patch("common.vault.base_vault_client.generate_bearer_token_from_creds", return_value=("sa_token_str", None))
+ @patch("common.vault.base_vault_client.log_info")
+ def test_get_bearer_token_generates_from_credentials_string(self, mock_log, mock_generate):
+ result = self.vault_client.get_bearer_token(CREDENTIALS_WITH_STRING)
+ mock_generate.assert_called_once()
+ self.assertEqual(result, "sa_token_str")
+
+ @patch("common.vault.base_vault_client.generate_bearer_token", return_value=("new_token", None))
+ @patch("common.vault.base_vault_client.is_expired", return_value=True)
+ @patch("common.vault.base_vault_client.log_info")
+ def test_get_bearer_token_regenerates_on_expiry(self, mock_log, mock_is_expired, mock_generate):
+ """Expired token is regenerated silently — no exception raised."""
+ self.vault_client._bearer_token = "expired_token"
+ result = self.vault_client.get_bearer_token(CREDENTIALS_WITH_PATH)
+ mock_generate.assert_called_once()
+ self.assertEqual(result, "new_token")
+
+ @patch("common.vault.base_vault_client.generate_bearer_token")
+ @patch("common.vault.base_vault_client.is_expired", return_value=False)
+ @patch("common.vault.base_vault_client.log_info")
+ def test_get_bearer_token_reuses_valid_cached_token(self, mock_log, mock_is_expired, mock_generate):
+ """Valid cached token is reused without calling generate_bearer_token."""
+ self.vault_client._bearer_token = "valid_token"
+ result = self.vault_client.get_bearer_token(CREDENTIALS_WITH_PATH)
+ mock_generate.assert_not_called()
+ self.assertEqual(result, "valid_token")
+
+ # ------------------------------------------------------------------ #
+ # update_config
+ # ------------------------------------------------------------------ #
+
+ def test_update_config_sets_flag(self):
+ self.vault_client.update_config({"credentials": "new_credentials"})
+ self.assertTrue(self.vault_client._is_config_updated)
+ self.assertEqual(self.vault_client.get_config()["credentials"], "new_credentials")
+
+ # ------------------------------------------------------------------ #
+ # get_current_bearer_token (new accessor -- did not exist pre-split)
+ # ------------------------------------------------------------------ #
+
+ def test_get_current_bearer_token_none_before_first_fetch(self):
+ self.assertIsNone(self.vault_client.get_current_bearer_token())
+
+ def test_get_current_bearer_token_returns_cached_value(self):
+ self.vault_client._bearer_token = "cached_value"
+ self.assertEqual(self.vault_client.get_current_bearer_token(), "cached_value")
+
+ # ------------------------------------------------------------------ #
+ # resolve_vault_url is abstract -- a subclass that forgets it can't instantiate
+ # ------------------------------------------------------------------ #
+
+ def test_resolve_vault_url_is_a_required_abstract_hook(self):
+ class MissingResolveVaultUrl(BaseVaultClient):
+ def initialize_api_client(self, vault_url, bearer_token):
+ pass
+
+ with self.assertRaises(TypeError):
+ MissingResolveVaultUrl(dict(CONFIG))
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/common/tests/vault/test_base_vault_controller.py b/common/tests/vault/test_base_vault_controller.py
new file mode 100644
index 00000000..8f0509cb
--- /dev/null
+++ b/common/tests/vault/test_base_vault_controller.py
@@ -0,0 +1,122 @@
+import unittest
+
+from common.errors import SkyflowError
+from common.utils import SkyflowMessages
+from common.vault.base_vault_controller import BaseVaultController
+
+
+class DummyVaultController(BaseVaultController):
+ _skyflow_messages = SkyflowMessages
+
+ def insert(self, request):
+ raise NotImplementedError
+
+ def get(self, request):
+ raise NotImplementedError
+
+ def update(self, request):
+ raise NotImplementedError
+
+ def delete(self, request):
+ raise NotImplementedError
+
+ def query(self, request):
+ raise NotImplementedError
+
+ def detokenize(self, request):
+ raise NotImplementedError
+
+
+class TestBaseVaultControllerAbstractContract(unittest.TestCase):
+ def test_cannot_instantiate_without_insert(self):
+ class Incomplete(BaseVaultController):
+ pass
+
+ with self.assertRaises(TypeError):
+ Incomplete(vault_client=None)
+
+ def test_cannot_instantiate_missing_any_single_method(self):
+ """Java-interface-style: every one of the six operations is independently required --
+ omitting any single one (not just insert) blocks instantiation."""
+ for missing in ("insert", "get", "update", "delete", "query", "detokenize"):
+ methods = {name: (lambda self, request: None) for name in
+ ("insert", "get", "update", "delete", "query", "detokenize") if name != missing}
+ Incomplete = type("Incomplete", (BaseVaultController,), methods)
+ with self.assertRaises(TypeError, msg=f"missing only '{missing}' should still fail to instantiate"):
+ Incomplete(vault_client=None)
+
+ def test_concrete_subclass_instantiates(self):
+ vault = DummyVaultController(vault_client=None)
+ self.assertIsInstance(vault, BaseVaultController)
+
+
+class TestValidateTableNameIfPresent(unittest.TestCase):
+ """Shared rule used identically by both variants (see PdbVaultController/flowvault's
+ VaultController.insert()): a table value, IF given, must be a non-empty string. Whether
+ table is required at all is variant-specific and stays out of this helper."""
+
+ def setUp(self):
+ self.vault = DummyVaultController(vault_client=None)
+
+ def test_none_is_allowed(self):
+ self.vault._validate_table_name_if_present(None) # should not raise
+
+ def test_valid_string_is_allowed(self):
+ self.vault._validate_table_name_if_present("table1") # should not raise
+
+ def test_empty_string_raises(self):
+ with self.assertRaises(SkyflowError):
+ self.vault._validate_table_name_if_present("")
+
+ def test_whitespace_only_string_raises(self):
+ with self.assertRaises(SkyflowError):
+ self.vault._validate_table_name_if_present(" ")
+
+ def test_non_string_raises(self):
+ with self.assertRaises(SkyflowError):
+ self.vault._validate_table_name_if_present(123)
+
+
+class TestValidateFieldValues(unittest.TestCase):
+ """Shared rule: a record's field-value map must be a non-empty dict of non-empty string
+ keys -- values may be anything, including None/empty string."""
+
+ def setUp(self):
+ self.vault = DummyVaultController(vault_client=None)
+
+ def test_valid_values_pass(self):
+ self.vault._validate_field_values({"name": "John", "age": 30}) # should not raise
+
+ def test_none_raises(self):
+ with self.assertRaises(SkyflowError):
+ self.vault._validate_field_values(None)
+
+ def test_non_dict_raises(self):
+ with self.assertRaises(SkyflowError):
+ self.vault._validate_field_values(["not", "a", "dict"])
+
+ def test_empty_dict_raises(self):
+ with self.assertRaises(SkyflowError):
+ self.vault._validate_field_values({})
+
+ def test_empty_key_raises(self):
+ with self.assertRaises(SkyflowError):
+ self.vault._validate_field_values({"": "value"})
+
+ def test_whitespace_only_key_raises(self):
+ with self.assertRaises(SkyflowError):
+ self.vault._validate_field_values({" ": "value"})
+
+ def test_none_value_is_valid(self):
+ self.vault._validate_field_values({"a": None}) # should not raise
+
+ def test_empty_string_value_is_valid(self):
+ self.vault._validate_field_values({"a": ""}) # should not raise
+
+ def test_falsy_non_string_values_are_valid(self):
+ """0, False, [], {} are all legitimate values."""
+ self.vault._validate_field_values({"a": 0, "b": False, "c": [], "d": {}}) # should not raise
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/common/utils/__init__.py b/common/utils/__init__.py
new file mode 100644
index 00000000..3dbcd1db
--- /dev/null
+++ b/common/utils/__init__.py
@@ -0,0 +1,4 @@
+from .enums import LogLevel, Env, TokenType
+from ._skyflow_messages import SkyflowMessages
+from ._helpers import get_base_url, format_scope, is_valid_url
+from ._utils import get_credentials, get_vault_url, validate_api_key
diff --git a/skyflow/utils/_helpers.py b/common/utils/_helpers.py
similarity index 100%
rename from skyflow/utils/_helpers.py
rename to common/utils/_helpers.py
diff --git a/common/utils/_skyflow_messages.py b/common/utils/_skyflow_messages.py
new file mode 100644
index 00000000..30f9d533
--- /dev/null
+++ b/common/utils/_skyflow_messages.py
@@ -0,0 +1,443 @@
+import sys
+from enum import Enum
+
+# `common` has no SDK_VERSION of its own -- picks up whichever variant is installed. Reads
+# sys.modules (never forces a fresh import) to avoid triggering a circular import back into
+# common.* mid-init.
+_version_module = sys.modules.get("skyflow.utils._version")
+SDK_VERSION = getattr(_version_module, "SDK_VERSION", "0.0.0")
+
+error_prefix = f"Skyflow Python SDK {SDK_VERSION}"
+INFO = "INFO"
+WARN = "WARN"
+ERROR = "ERROR"
+
+class SkyflowMessages:
+ class ErrorCodes(Enum):
+ INVALID_INPUT = 400
+ INVALID_INDEX = 404
+ SERVER_ERROR = 500
+ PARTIAL_SUCCESS = 500
+ TOKENS_GET_COLUMN_NOT_SUPPORTED = 400
+ REDACTION_WITH_TOKENS_NOT_SUPPORTED = 400
+
+ class Error(Enum):
+ GENERIC_API_ERROR = f"{error_prefix} API error. Error occurred."
+
+ EMPTY_VAULT_ID = f"{error_prefix} Initialization failed. Invalid vault Id. Specify a valid vault Id."
+ INVALID_VAULT_ID = f"{error_prefix} Initialization failed. Invalid vault Id. Specify a valid vault Id as a string."
+ EMPTY_CLUSTER_ID = f"{error_prefix} Initialization failed. Invalid cluster Id for vault with id {{}}. Specify a valid cluster Id."
+ INVALID_CLUSTER_ID = f"{error_prefix} Initialization failed. Invalid cluster Id for vault with id {{}}. Specify cluster Id as a string."
+ INVALID_ENV = f"{error_prefix} Initialization failed. Invalid env for vault with id {{}}. Specify a valid env."
+ INVALID_KEY = f"{error_prefix} Initialization failed. Invalid {{}}. Specify a valid key"
+ VAULT_ID_NOT_IN_CONFIG_LIST = f"{error_prefix} Validation error. Vault id {{}} is missing from the config. Specify the vault id from configs."
+ EMPTY_VAULT_CONFIGS = f"{error_prefix} Validation error. Specify at least one vault config."
+ EMPTY_CONNECTION_CONFIGS = f"{error_prefix} Validation error. Specify at least one connection config."
+ VAULT_ID_ALREADY_EXISTS =f"{error_prefix} Initialization failed. vault with id {{}} already exists."
+ CONNECTION_ID_ALREADY_EXISTS = f"{error_prefix} Initialization failed. Connection with id {{}} already exists."
+
+ EMPTY_CREDENTIALS = f"{error_prefix} Validation error. Invalid credentials for {{}} with id {{}}. Credentials must not be empty."
+ INVALID_CREDENTIALS_IN_CONFIG = f"{error_prefix} Validation error. Invalid credentials for {{}} with id {{}}. Specify a valid credentials."
+ INVALID_CREDENTIALS = f"{error_prefix} Validation error. Invalid credentials. Specify a valid credentials."
+ MULTIPLE_CREDENTIALS_PASSED_IN_CONFIG = f"{error_prefix} Validation error. Multiple credentials provided for {{}} with id {{}}. Please specify only one valid credential."
+ MULTIPLE_CREDENTIALS_PASSED = f"{error_prefix} Validation error. Multiple credentials provided. Please specify only one valid credential."
+ EMPTY_CREDENTIALS_STRING_IN_CONFIG = f"{error_prefix} Validation error. Invalid credentials for {{}} with id {{}}. Specify valid credentials."
+ EMPTY_CREDENTIALS_STRING = f"{error_prefix} Validation error. Invalid credentials. Specify valid credentials."
+ INVALID_CREDENTIALS_STRING_IN_CONFIG = f"{error_prefix} Validation error. Invalid credentials for {{}} with id {{}}. Specify credentials as a string."
+ INVALID_CREDENTIALS_STRING = f"{error_prefix} Validation error. Invalid credentials. Specify credentials as a string."
+ EMPTY_CREDENTIAL_FILE_PATH_IN_CONFIG = f"{error_prefix} Initialization failed. Invalid credentials for {{}} with id {{}}. Specify a valid file path."
+ EMPTY_CREDENTIAL_FILE_PATH = f"{error_prefix} Initialization failed. Invalid credentials. Specify a valid file path."
+ INVALID_CREDENTIAL_FILE_PATH_IN_CONFIG = f"{error_prefix} Initialization failed. Invalid credentials for {{}} with id {{}}. Expected file path to be a string."
+ INVALID_CREDENTIAL_FILE_PATH = f"{error_prefix} Initialization failed. Invalid credentials. Expected file path to be a valid file path."
+ EMPTY_CREDENTIALS_TOKEN_IN_CONFIG = f"{error_prefix} Initialization failed. Invalid token for {{}} with id {{}}.Specify a valid credentials token."
+ EMPTY_CREDENTIALS_TOKEN = f"{error_prefix} Initialization failed. Invalid token.Specify a valid credentials token."
+ INVALID_CREDENTIALS_TOKEN_IN_CONFIG = f"{error_prefix} Initialization failed. Invalid credentials token for {{}} with id {{}}. Expected token to be a string."
+ INVALID_CREDENTIALS_TOKEN = f"{error_prefix} Initialization failed. Invalid credentials token. Expected token to be a string."
+ EXPIRED_BEARER_TOKEN = f"{error_prefix} Initialization failed. Bearer token is invalid or expired."
+ EXPIRED_TOKEN = f"{error_prefix} Initialization failed. Given token is expired. Specify a valid credentials token."
+ EMPTY_API_KEY_IN_CONFIG = f"{error_prefix} Initialization failed. Invalid api key for {{}} with id {{}}.Specify a valid api key."
+ EMPTY_API_KEY= f"{error_prefix} Initialization failed. Invalid api key.Specify a valid api key."
+ INVALID_API_KEY_IN_CONFIG = f"{error_prefix} Initialization failed. Invalid api key for {{}} with id {{}}. Expected api key to be a string."
+ INVALID_API_KEY = f"{error_prefix} Initialization failed. Invalid api key. Expected api key to be a string."
+ INVALID_ROLES_KEY_TYPE_IN_CONFIG = f"{error_prefix} Validation error. Invalid roles for {{}} with id {{}}. Specify roles as an array."
+ INVALID_ROLES_KEY_TYPE = f"{error_prefix} Validation error. Invalid roles. Specify roles as an array."
+ EMPTY_ROLES_IN_CONFIG = f"{error_prefix} Validation error. Invalid roles for {{}} with id {{}}. Specify at least one role."
+ EMPTY_ROLES = f"{error_prefix} Validation error. Invalid roles. Specify at least one role."
+ EMPTY_CONTEXT_IN_CONFIG = f"{error_prefix} Initialization failed. Invalid context provided for {{}} with id {{}}. Specify context as type Context."
+ EMPTY_CONTEXT = f"{error_prefix} Initialization failed. Invalid context provided. Specify context as type Context."
+ INVALID_CONTEXT_IN_CONFIG = f"{error_prefix} Initialization failed. Invalid context for {{}} with id {{}}. Specify a valid context."
+ INVALID_CONTEXT = f"{error_prefix} Initialization failed. Invalid context. Specify a valid context."
+ INVALID_CTX_TYPE = f"{error_prefix} Initialization failed. Invalid ctx type. Specify ctx as a string or a dict."
+ INVALID_CTX_MAP_KEY = f"{error_prefix} Initialization failed. Invalid key '{{}}' in ctx dict. Keys must contain only alphanumeric characters and underscores."
+ INVALID_LOG_LEVEL = f"{error_prefix} Initialization failed. Invalid log level. Specify a valid log level."
+ EMPTY_LOG_LEVEL = f"{error_prefix} Initialization failed. Specify a valid log level."
+
+ EMPTY_CONNECTION_ID = f"{error_prefix} Initialization failed. Invalid connection Id. Specify a valid connection Id."
+ INVALID_CONNECTION_ID = f"{error_prefix} Initialization failed. Invalid connection Id. Specify connection Id as a string."
+ EMPTY_CONNECTION_URL = f"{error_prefix} Initialization failed. Invalid connection Url for connection with id {{}}. Specify a valid connection Url."
+ INVALID_CONNECTION_URL = f"{error_prefix} Initialization failed. Invalid connection Url for connection with id {{}}. Specify connection Url as a string."
+ CONNECTION_ID_NOT_IN_CONFIG_LIST = f"{error_prefix} Validation error. {{}} is missing from the config. Specify the connectionIds from config."
+ RESPONSE_NOT_JSON = f"{error_prefix} Response {{}} is not valid JSON."
+ API_ERROR = f"{error_prefix} Server returned status code {{}}"
+
+ INVALID_JSON_RESPONSE = f"{error_prefix} Invalid JSON response received."
+ UNKNOWN_ERROR_DEFAULT_MESSAGE = f"{error_prefix} An unknown error occurred."
+
+ INVALID_FILE_INPUT = f"{error_prefix} Validation error. Invalid file input. Specify a valid file input."
+ INVALID_DETECT_ENTITIES_TYPE = f"{error_prefix} Validation error. Invalid type of detect entities. Specify detect entities as list of DetectEntities enum."
+ INVALID_TYPE_FOR_DEFAULT_TOKEN_TYPE = f"{error_prefix} Validation error. Invalid type of default token type. Specify default token type as TokenType enum."
+ INVALID_TOKEN_TYPE_VALUE = f"{error_prefix} Validation error. Invalid value for token type {{}}. Specify as list of DetectEntities enum."
+ INVALID_MAXIMUM_RESOLUTION = f"{error_prefix} Validation error. Invalid type of maximum resolution. Specify maximum resolution as a number."
+ INVALID_OUTPUT_DIRECTORY_VALUE = f"{error_prefix} Validation error. Invalid type of output directory. Specify output directory as a string."
+ WAIT_TIME_GREATER_THEN_64 = f"{error_prefix} Validation error. Invalid wait time. The waitTime value must be between 0 and 64 seconds."
+ OUTPUT_DIRECTORY_NOT_FOUND = f"{error_prefix} Validation error. Invalid output directory. Directory {{}} not found."
+
+ MISSING_TABLE_NAME_IN_INSERT = f"{error_prefix} Validation error. Table name cannot be empty in insert request. Specify a table name."
+ INVALID_TABLE_NAME_IN_INSERT = f"{error_prefix} Validation error. Invalid table name in insert request. Specify a valid table name."
+ INVALID_TYPE_OF_DATA_IN_INSERT = f"{error_prefix} Validation error. Invalid type of data in insert request. Specify data as a object array."
+ EMPTY_DATA_IN_INSERT = f"{error_prefix} Validation error. Data array cannot be empty. Specify data in insert request."
+ INVALID_RECORD_DATA_IN_INSERT = f"{error_prefix} Validation error. Each record's field values must be a non-empty dict."
+ EMPTY_KEY_IN_INSERT_DATA = f"{error_prefix} Validation error. A record must not contain a null or empty key."
+ EMPTY_VALUE_IN_INSERT_DATA = f"{error_prefix} Validation error. A record must not contain a null or empty value."
+ INVALID_UPSERT_OPTIONS_TYPE = f"{error_prefix} Validation error. Invalid 'upsert' value in options. Specify 'upsert' as a non-empty string containing the column name."
+ INVALID_HOMOGENEOUS_TYPE = f"{error_prefix} Validation error. Invalid type of homogeneous. Specify homogeneous as a string."
+ INVALID_TOKEN_MODE_TYPE = f"{error_prefix} Validation error. Invalid type of token mode. Specify token mode as a TokenMode enum."
+ INVALID_RETURN_TOKENS_TYPE = f"{error_prefix} Validation error. Invalid type of return tokens. Specify return tokens as a boolean."
+ INVALID_CONTINUE_ON_ERROR_TYPE = f"{error_prefix} Validation error. Invalid type of continue on error. Specify continue on error as a boolean."
+ TOKENS_PASSED_FOR_TOKEN_MODE_DISABLE = f"{error_prefix} Validation error. 'token_mode' wasn't specified. Set 'token_mode' to 'ENABLE' to insert tokens."
+ INSUFFICIENT_TOKENS_PASSED_FOR_TOKEN_MODE_ENABLE_STRICT = f"{error_prefix} Validation error. 'token_mode' is set to 'ENABLE_STRICT', but some fields are missing tokens. Specify tokens for all fields."
+ MISMATCH_OF_FIELDS_AND_TOKENS = f"{error_prefix} Validation error. Keys for values and tokens are not matching. Ensure each values entry and its corresponding tokens entry have the same keys."
+ NO_TOKENS_IN_INSERT = f"{error_prefix} Validation error. Tokens weren't specified for records while 'token_mode' was {{}}. Specify tokens."
+ BATCH_INSERT_FAILURE = f"{error_prefix} Insert operation failed."
+ GET_FAILURE = f"{error_prefix} Get operation failed."
+ HOMOGENOUS_NOT_SUPPORTED_WITH_UPSERT = f"{error_prefix} Validation error. Homogenous is not supported when upsert is passed."
+
+ EMPTY_TABLE_VALUE = f"{error_prefix} Validation error. 'table' can't be empty. Specify a table."
+ INVALID_TABLE_VALUE = f"{error_prefix} Validation error. Invalid type of table. Specify table as a string"
+ EMPTY_RECORD_IDS_IN_DELETE = f"{error_prefix} Validation error. 'record ids' array can't be empty. Specify one or more record ids."
+ BULK_DELETE_FAILURE = f"{error_prefix} Delete operation failed."
+ EMPTY_SKYFLOW_ID= f"{error_prefix} Validation error. skyflow_id can't be empty."
+ INVALID_FILE_COLUMN_NAME= f"{error_prefix} Validation error. 'column_name' can't be empty."
+
+ INVALID_QUERY_TYPE = f"{error_prefix} Validation error. Query parameter is of type {{}}. Specify as a string."
+ EMPTY_QUERY = f"{error_prefix} Validation error. Query parameter can't be empty. Specify as a string."
+ INVALID_QUERY_COMMAND = f"{error_prefix} Validation error. {{}} command was passed instead, but only SELECT commands are supported. Specify the SELECT command."
+ SERVER_ERROR = f"{error_prefix} Validation error. Check SkyflowError.data for details."
+ QUERY_FAILED = f"{error_prefix} Query operation failed."
+ DETOKENIZE_FIELD = f"{error_prefix} Detokenize operation failed."
+ UPDATE_FAILED = f"{error_prefix} Update operation failed."
+ TOKENIZE_FAILED = f"{error_prefix} Tokenize operation failed."
+ INVOKE_CONNECTION_FAILED = f"{error_prefix} Invoke Connection operation failed."
+
+ INVALID_IDS_TYPE = f"{error_prefix} Validation error. 'ids' has a value of type {{}}. Specify 'ids' as list."
+ INVALID_REDACTION_TYPE = f"{error_prefix} Validation error. 'redaction_type' has a value of type {{}}. Specify 'redaction_type' as type Skyflow.RedactionType."
+ INVALID_COLUMN_NAME = f"{error_prefix} Validation error. column_name has a value of type {{}}. Specify 'column' as a string."
+ INVALID_COLUMN_VALUE = f"{error_prefix} Validation error. column_values key has a value of type {{}}. Specify column_values key as list."
+ INVALID_COLUMN_VALUES = f"{error_prefix} Validation error. column_values key is an empty list. Specify at least one column value when column_name is passed."
+ INVALID_FIELDS_VALUE = f"{error_prefix} Validation error. fields key has a value of type{{}}. Specify fields key as list."
+ BOTH_OFFSET_AND_LIMIT_SPECIFIED = f"{error_prefix} Validation error. Both offset and limit cannot be present at the same time"
+ INVALID_OFF_SET_VALUE = f"{error_prefix} Validation error. offset key has a value of type {{}}. Specify offset key as integer."
+ INVALID_LIMIT_VALUE = f"{error_prefix} Validation error. limit key has a value of type {{}}. Specify limit key as integer."
+ INVALID_DOWNLOAD_URL_VALUE = f"{error_prefix} Validation error. download_url key has a value of type {{}}. Specify download_url key as boolean."
+ REDACTION_WITH_TOKENS_NOT_SUPPORTED = f"{error_prefix} Validation error. 'redaction_type' can't be used when tokens are specified. Remove 'redaction_type' from payload if tokens are specified."
+ TOKENS_GET_COLUMN_NOT_SUPPORTED = f"{error_prefix} Validation error. Column name and/or column values can't be used when tokens are specified. Remove unique column values or tokens from the payload."
+ BOTH_IDS_AND_COLUMN_DETAILS_SPECIFIED = f"{error_prefix} Validation error. Both Skyflow IDs and column details can't be specified. Either specify Skyflow IDs or unique column details."
+ INVALID_ORDER_BY_VALUE = f"{error_prefix} Validation error. order_by key has a value of type {{}}. Specify order_by key as Skyflow.OrderBy"
+
+ UPDATE_FIELD_KEY_ERROR = f"{error_prefix} Validation error. Fields are empty in an update payload. Specify at least one field."
+ INVALID_FIELDS_TYPE = f"{error_prefix} Validation error. The 'data' key has a value of type {{}}. Specify 'data' as a dictionary."
+ IDS_KEY_ERROR = f"{error_prefix} Validation error. 'ids' key is missing from the payload. Specify an 'ids' key."
+ INVALID_TOKENS_LIST_VALUE = f"{error_prefix} Validation error. The 'data' field is invalid. Specify 'data' as a list of dictionaries containing 'token' and 'redaction_type'."
+ INVALID_DATA_FOR_DETOKENIZE = f"{error_prefix}"
+ EMPTY_TOKENS_LIST_VALUE = f"{error_prefix} Validation error. Tokens are empty in detokenize payload. Specify at lease one token"
+ INVALID_TOKEN_TYPE = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Tokens should be of type string."
+
+ INVALID_TOKENIZE_PARAMETERS = f"{error_prefix} Validation error. The 'values' key has a value of type {{}}. Specify 'tokenize_parameters' as a list."
+ EMPTY_TOKENIZE_PARAMETERS = f"{error_prefix} Validation error. Tokenize values are empty in tokenize payload. Specify at least one parameter."
+ INVALID_TOKENIZE_PARAMETER = f"{error_prefix} Validation error. Tokenize value at index {{}} has a value of type {{}}. Specify as a dictionary."
+ EMPTY_TOKENIZE_PARAMETER_VALUE = f"{error_prefix} Validation error. Tokenize value at index {{}} is empty. Specify a valid value."
+ EMPTY_TOKENIZE_PARAMETER_COLUMN_GROUP = f"{error_prefix} Validation error. Tokenize column group at index {{}} is empty. Specify a valid column group."
+ INVALID_TOKENIZE_PARAMETER_KEY = f"{error_prefix} Validation error. Tokenize value key at index {{}} is invalid. Specify a valid key value."
+
+ INVALID_REQUEST_BODY = f"{error_prefix} Validation error. Invalid request body. Specify the request body as an object."
+ INVALID_REQUEST_HEADERS = f"{error_prefix} Validation error. Invalid request headers. Specify the request as an object."
+ INVALID_URL = f"{error_prefix} Validation error. Connection url {{}} is invalid. Specify a valid connection url."
+ INVALID_PATH_PARAMS = f"{error_prefix} Validation error. Path parameters aren't valid. Specify valid path parameters."
+ INVALID_QUERY_PARAMS = f"{error_prefix} Validation error. Query parameters aren't valid. Specify valid query parameters."
+ INVALID_REQUEST_METHOD = f"{error_prefix} Validation error. Invalid request method. Specify the request method as enum RequestMethod"
+
+ MISSING_PRIVATE_KEY = f"{error_prefix} Initialization failed. Unable to read private key in credentials. Verify your private key."
+ MISSING_CLIENT_ID = f"{error_prefix} Initialization failed. Unable to read client ID in credentials. Verify your client ID."
+ MISSING_KEY_ID = f"{error_prefix} Initialization failed. Unable to read key ID in credentials. Verify your key ID."
+ MISSING_TOKEN_URI = f"{error_prefix} Initialization failed. Unable to read token URI in credentials. Verify your token URI."
+ INVALID_TOKEN_URI = f"{error_prefix} Initialization failed. Invalid Skyflow credentials. The token URI must be a string and a valid URL."
+ JWT_INVALID_FORMAT = f"{error_prefix} Initialization failed. Invalid private key format. Verify your credentials."
+ JWT_DECODE_ERROR = f"{error_prefix} Validation error. Invalid access token. Verify your credentials."
+ FILE_INVALID_JSON = f"{error_prefix} Initialization failed. File at {{}} is not in valid JSON format. Verify the file contents."
+ INVALID_JSON_FORMAT_IN_CREDENTIALS_ENV = f"{error_prefix} Validation error. Invalid JSON format in SKYFLOW_CREDENTIALS environment variable."
+ FAILED_TO_GET_BEARER_TOKEN = f"{ERROR}: [{error_prefix}] Failed to generate bearer token."
+ UNAUTHORIZED_ERROR_IN_GETTING_BEARER_TOKEN = f"{ERROR}: [{error_prefix}] Authorization failed while retrieving the bearer token."
+
+ INVALID_TEXT_IN_DEIDENTIFY= f"{error_prefix} Validation error. The text field is required and must be a non-empty string. Specify a valid text."
+ INVALID_ENTITIES_IN_DEIDENTIFY= f"{error_prefix} Validation error. The entities field must be an array of DetectEntities enums. Specify a valid entities."
+ INVALID_ALLOW_REGEX_LIST= f"{error_prefix} Validation error. The allowRegexList field must be an array of strings. Specify a valid allow_regex_list."
+ INVALID_RESTRICT_REGEX_LIST= f"{error_prefix} Validation error. The restrictRegexList field must be an array of strings. Specify a valid restrict_regex_list."
+ INVALID_TOKEN_FORMAT= f"{error_prefix} Validation error. The tokenFormat key must be an instance of TokenFormat. Specify a valid token format."
+ INVALID_TRANSFORMATIONS= f"{error_prefix} Validation error. The transformations key must be an instance of Transformations. Specify a valid transformations."
+
+ INVALID_TEXT_IN_REIDENTIFY= f"{error_prefix} Validation error. The text field is required and must be a non-empty string. Specify a valid text."
+ INVALID_REDACTED_ENTITIES_IN_REIDENTIFY= f"{error_prefix} Validation error. The redactedEntities field must be an array of DetectEntities enums. Specify a valid redactedEntities."
+ INVALID_MASKED_ENTITIES_IN_REIDENTIFY= f"{error_prefix} Validation error. The maskedEntities field must be an array of DetectEntities enums. Specify a valid maskedEntities."
+ INVALID_PLAIN_TEXT_ENTITIES_IN_REIDENTIFY= f"{error_prefix} Validation error. The plainTextEntities field must be an array of DetectEntities enums. Specify a valid plainTextEntities."
+
+ INVALID_DEIDENTIFY_FILE_REQUEST= f"{error_prefix} Validation error. Invalid deidentify file request. Specify a valid deidentify file request."
+ INVALID_DEIDENTIFY_FILE_INPUT= f"{error_prefix} Validation error. Invalid deidentify file input. Please provide either a file or a file path."
+ EMPTY_FILE_OBJECT= f"{error_prefix} Validation error. File object cannot be empty. Specify a valid file object."
+ INVALID_FILE_FORMAT= f"{error_prefix} Validation error. Invalid file format. Specify a valid file format."
+ MISSING_FILE_SOURCE= f"{error_prefix} Validation error. Provide exactly one of filePath, base64, or fileObject."
+ INVALID_FILE_OBJECT= f"{error_prefix} Validation error. Invalid file object. Specify a valid file object."
+ INVALID_BASE64_STRING= f"{error_prefix} Validation error. Invalid base64 string. Specify a valid base64 string."
+ INVALID_DEIDENTIFY_FILE_OPTIONS= f"{error_prefix} Validation error. Invalid deidentify file options. Specify a valid deidentify file options."
+ INVALID_ENTITIES= f"{error_prefix} Validation error. Invalid entities. Specify valid entities as string array."
+ EMPTY_ENTITIES= f"{error_prefix} Validation error. Entities cannot be empty. Specify valid entities."
+ EMPTY_ALLOW_REGEX_LIST= f"{error_prefix} Validation error. Allow regex list cannot be empty. Specify valid allow regex list."
+ INVALID_ALLOW_REGEX= f"{error_prefix} Validation error. Invalid allow regex. Specify valid allow regex at index {{}}."
+ EMPTY_RESTRICT_REGEX_LIST= f"{error_prefix} Validation error. Restrict regex list cannot be empty. Specify valid restrict regex list."
+ INVALID_RESTRICT_REGEX= f"{error_prefix} Validation error. Invalid restrict regex. Specify valid restrict regex at index {{}}."
+ INVALID_OUTPUT_PROCESSED_IMAGE= f"{error_prefix} Validation error. Invalid output processed image. Specify valid output processed image as boolean."
+ INVALID_OUTPUT_OCR_TEXT= f"{error_prefix} Validation error. Invalid output ocr text. Specify valid output ocr text as boolean."
+ INVALID_MASKING_METHOD= f"{error_prefix} Validation error. Invalid masking method. Specify valid masking method as MaskingMethod enum."
+ INVALID_PIXEL_DENSITY= f"{error_prefix} Validation error. Invalid pixel density. Specify valid pixel density as number."
+ INVALID_OUTPUT_TRANSCRIPTION= f"{error_prefix} Validation error. Invalid output transcription. Specify valid output transcription as DetectOutputTranscriptions enum."
+ INVALID_BLEEP_TYPE= f"{error_prefix} Validation error. Invalid type of bleep. Specify bleep as Bleep object."
+ INVALID_BLEEP_GAIN= f"{error_prefix} Validation error. Invalid bleep gain. Specify valid bleep gain as a number."
+ INVALID_BLEEP_FREQUENCY= f"{error_prefix} Validation error. Invalid bleep frequency. Specify valid bleep frequency as a number."
+ INVALID_BLEEP_START_PADDING= f"{error_prefix} Validation error. Invalid bleep start padding. Specify valid bleep start padding as a number."
+ INVALID_BLEEP_STOP_PADDING= f"{error_prefix} Validation error. Invalid bleep stop padding. Specify valid bleep stop padding as a number."
+ INVALID_OUTPUT_PROCESSED_AUDIO= f"{error_prefix} Validation error. Invalid output processed audio. Specify valid output processed audio as boolean."
+ INVALID_MAX_RESOLUTION= f"{error_prefix} Validation error. Invalid max resolution. Specify valid max resolution as string."
+ INVALID_BLEEP= f"{error_prefix} Validation error. Invalid bleep. Specify valid bleep as object."
+ INVALID_FILE_OR_ENCODED_FILE= f"{error_prefix} . Error while decoding base64 and saving file"
+ INVALID_FILE_TYPE = f"{error_prefix} Validation error. Invalid file type. Specify a valid file type."
+ INVALID_FILE_NAME= f"{error_prefix} Validation error. Invalid file name. Specify a valid file name."
+ INVALID_FILE_PATH= f"{error_prefix} Validation error. Invalid file path. Specify a valid file path."
+ INVALID_DEIDENTIFY_FILE_PATH= f"{error_prefix} Validation error. Invalid file path. Specify a valid file path."
+ INVALID_BASE64_HEADER= f"{error_prefix} Validation error. Invalid base64 header. Specify a valid base64 header."
+ INVALID_WAIT_TIME= f"{error_prefix} Validation error. Invalid wait time. Specify a valid wait time as number and should not be greater than 64 secs."
+ INVALID_OUTPUT_DIRECTORY= f"{error_prefix} Validation error. Invalid output directory. Specify a valid output directory as string."
+ INVALID_OUTPUT_DIRECTORY_PATH= f"{error_prefix} Validation error. Invalid output directory path. Specify a valid output directory path as string."
+ EMPTY_RUN_ID= f"{error_prefix} Validation error. Run id cannot be empty. Specify a valid run id."
+ INVALID_RUN_ID= f"{error_prefix} Validation error. Invalid run id. Specify a valid run id as string."
+ INTERNAL_SERVER_ERROR= f"{error_prefix}. Internal server error. {{}}."
+ GET_DETECT_RUN_FAILED = f"{error_prefix} Get detect run operation failed."
+ BASE_SKYFLOW_INSTANTIATION_NOT_ALLOWED = f"{error_prefix} BaseSkyflow cannot be instantiated directly. Build a concrete Skyflow class via make_skyflow_class() instead."
+
+ class Info(Enum):
+ CLIENT_INITIALIZED = f"{INFO}: [{error_prefix}] Initialized skyflow client."
+ VALIDATING_VAULT_CONFIG = f"{INFO}: [{error_prefix}] Validating vault config."
+ VALIDATING_CONNECTION_CONFIG = f"{INFO}: [{error_prefix}] Validating connection config."
+ UNABLE_TO_GENERATE_SDK_METRIC = f"{INFO}: [{error_prefix}] Unable to generate {{}} metric."
+ VAULT_CONTROLLER_INITIALIZED = f"{INFO}: [{error_prefix}] Initialized vault controller with vault ID {{}}."
+ CONNECTION_CONTROLLER_INITIALIZED = f"{INFO}: [{error_prefix}] Initialized connection controller with connection ID {{}}."
+ DETECT_CONTROLLER_INITIALIZED = f"{INFO}: [{error_prefix}] Initialized detect controller with vault ID {{}}."
+ VAULT_CONFIG_EXISTS = f"{INFO}: [{error_prefix}] Vault config with vault ID {{}} already exists."
+ VAULT_CONFIG_DOES_NOT_EXIST = f"{INFO}: [{error_prefix}] Vault config with vault ID {{}} doesn't exist."
+ CONNECTION_CONFIG_EXISTS = f"{INFO}: [{error_prefix}] Connection config with connection ID {{}} already exists."
+ CONNECTION_CONFIG_DOES_NOT_EXIST = f"{INFO}: [{error_prefix}] Connection config with connection ID {{}} doesn't exist."
+ LOGGER_SETUP_DONE = f"{INFO}: [{error_prefix}] Set up logger."
+ CURRENT_LOG_LEVEL = f"{INFO}: [{error_prefix}] Current log level is {{}}."
+
+ BEARER_TOKEN_EXPIRED = f"{INFO}: [{error_prefix}] Bearer token is expired."
+ GET_BEARER_TOKEN_TRIGGERED = f"{INFO}: [{error_prefix}] generate_bearer_token method triggered."
+ GET_BEARER_TOKEN_SUCCESS = f"{INFO}: [{error_prefix}] Bearer token generated."
+ GET_SIGNED_DATA_TOKENS_TRIGGERED = f"{INFO}: [{error_prefix}] generate_signed_data_tokens method triggered."
+ GET_SIGNED_DATA_TOKEN_SUCCESS = f"{INFO}: [{error_prefix}] Signed data tokens generated."
+ GENERATE_BEARER_TOKEN_FROM_CREDENTIALS_STRING_TRIGGERED = f"{INFO}: [{error_prefix}] generate bearer_token_from_credential_string method triggered."
+ REUSE_BEARER_TOKEN = f"{INFO}: [{error_prefix}] Reusing bearer token."
+
+ VALIDATE_DEIDENTIFY_FILE_REQUEST = f"{INFO}: [{error_prefix}] Validating deidentify file request."
+ DETECT_FILE_TRIGGERED = f"{INFO}: [{error_prefix}] Detect file method triggered."
+ DETECT_FILE_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Deidentify file request resolved."
+ DETECT_FILE_SUCCESS = f"{INFO}: [{error_prefix}] File deidentified."
+
+ VALIDATE_INSERT_REQUEST = f"{INFO}: [{error_prefix}] Validating insert request."
+ INSERT_TRIGGERED = f"{INFO}: [{error_prefix}] Insert method triggered."
+ INSERT_SUCCESS = f"{INFO}: [{error_prefix}] Data inserted."
+ INSERT_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Insert request resolved."
+
+ VALIDATE_UPDATE_REQUEST = f"{INFO}: [{error_prefix}] Validating update request."
+ UPDATE_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Update request resolved."
+ UPDATE_SUCCESS = f"{INFO}: [{error_prefix}] Data updated."
+ UPDATE_TRIGGERED = f"{INFO}: [{error_prefix}] Update method triggered."
+
+ DELETE_TRIGGERED = f"{INFO}: [{error_prefix}] Delete method triggered."
+ VALIDATING_DELETE_REQUEST = f"{INFO}: [{error_prefix}] Validating delete request."
+ DELETE_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Delete request resolved."
+ DELETE_SUCCESS = f"{INFO}: [{error_prefix}] Data deleted."
+
+ GET_TRIGGERED = f"{INFO}: [{error_prefix}] Get method triggered."
+ VALIDATE_GET_REQUEST = f"{INFO}: [{error_prefix}] Validating get request."
+ GET_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Get request resolved."
+ GET_SUCCESS = f"{INFO}: [{error_prefix}] Data revealed."
+
+ QUERY_TRIGGERED = f"{INFO}: [{error_prefix}] Query method triggered."
+ VALIDATING_QUERY_REQUEST = f"{INFO}: [{error_prefix}] Validating query request."
+ QUERY_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Query request resolved."
+ QUERY_SUCCESS = f"{INFO}: [{error_prefix}] Query executed."
+
+ DETOKENIZE_TRIGGERED = f"{INFO}: [{error_prefix}] Detokenize method triggered."
+ VALIDATE_DETOKENIZE_REQUEST = f"{INFO}: [{error_prefix}] Validating detokenize request."
+ DETOKENIZE_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Detokenize request resolved."
+ DETOKENIZE_SUCCESS = f"{INFO}: [{error_prefix}] Data detokenized."
+
+ TOKENIZE_TRIGGERED = f"{INFO}: [{error_prefix}] Tokenize method triggered."
+ VALIDATING_TOKENIZE_REQUEST = f"{INFO}: [{error_prefix}] Validating tokenize request."
+ TOKENIZE_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Tokenize request resolved."
+ TOKENIZE_SUCCESS = f"{INFO}: [{error_prefix}] Data tokenized."
+
+ FILE_UPLOAD_TRIGGERED = f"{INFO}: [{error_prefix}] File upload method triggered."
+ VALIDATING_FILE_UPLOAD_REQUEST = f"{INFO}: [{error_prefix}] Validating file upload request."
+ FILE_UPLOAD_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] File upload request resolved."
+ FILE_UPLOAD_SUCCESS = f"{INFO}: [{error_prefix}] File uploaded successfully."
+
+ INVOKE_CONNECTION_TRIGGERED = f"{INFO}: [{error_prefix}] Invoke connection method triggered."
+ VALIDATING_INVOKE_CONNECTION_REQUEST = f"{INFO}: [{error_prefix}] Validating invoke connection request."
+ INVOKE_CONNECTION_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Invoke connection request resolved."
+ INVOKE_CONNECTION_SUCCESS = f"{INFO}: [{error_prefix}] Invoke Connection Success."
+
+ DEIDENTIFY_TEXT_TRIGGERED = f"{INFO}: [{error_prefix}] Deidentify text method triggered."
+ VALIDATING_DEIDENTIFY_TEXT_INPUT = f"{INFO}: [{error_prefix}] Validating deidentify text input."
+ DEIDENTIFY_TEXT_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Deidentify text request is resolved."
+ DEIDENTIFY_TEXT_SUCCESS = f"{INFO}: [{error_prefix}] Data deidentified."
+
+ REIDENTIFY_TEXT_TRIGGERED = f"{INFO}: [{error_prefix}] Reidentify text method triggered."
+ VALIDATING_REIDENTIFY_TEXT_INPUT = f"{INFO}: [{error_prefix}] Validating reidentify text input."
+ REIDENTIFY_TEXT_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Reidentify text request is resolved."
+ REIDENTIFY_TEXT_SUCCESS = f"{INFO}: [{error_prefix}] Data reidentified."
+
+ DEIDENTIFY_FILE_TRIGGERED = f"{INFO}: [{error_prefix}] Deidentify file triggered."
+ VALIDATING_DETECT_FILE_INPUT = f"{INFO}: [{error_prefix}] Validating deidentify file input."
+ DEIDENTIFY_FILE_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Deidentify file request is resolved."
+ DEIDENTIFY_FILE_SUCCESS = f"{INFO}: [{error_prefix}] File deidentified."
+
+ GET_DETECT_RUN_TRIGGERED = f"{INFO}: [{error_prefix}] Get detect run triggered."
+ VALIDATING_GET_DETECT_RUN_INPUT = f"{INFO}: [{error_prefix}] Validating get detect run input."
+ GET_DETECT_RUN_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Get detect run request is resolved."
+ GET_DETECT_RUN_SUCCESS = f"{INFO}: [{error_prefix}] Get detect run success."
+
+ DETECT_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Detect request is resolved."
+
+ class ErrorLogs(Enum):
+ INVALID_LOG_LEVEL = f"{ERROR}: [{error_prefix}] Invalid log level. Specify a valid log level."
+ INVALID_KEY = f"{ERROR}: [{error_prefix}] Invalid key {{}} in config."
+ VAULTID_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Invalid vault config. Vault ID is required."
+ EMPTY_VAULTID = f"{ERROR}: [{error_prefix}] Invalid vault config. Vault ID can not be empty."
+ CLUSTER_ID_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Invalid vault config. Cluster ID is required."
+ EMPTY_CLUSTER_ID = f"{ERROR}: [{error_prefix}] Invalid vault config. Cluster ID can not be empty."
+ ENV_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Invalid vault config. Env is required."
+ CONNECTION_ID_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Invalid connection config. Connection ID is required."
+ EMPTY_CONNECTION_ID = f"{ERROR}: [{error_prefix}] Invalid connection config. Connection ID can not be empty."
+ CONNECTION_URL_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Invalid connection config. Connection URL is required."
+ EMPTY_CONNECTION_URL = f"{ERROR}: [{error_prefix}] Invalid connection config. Connection URL can not be empty."
+ INVALID_CONNECTION_URL = f"{ERROR}: [{error_prefix}] Invalid connection config. Connection URL is not a valid URL."
+ EMPTY_CREDENTIALS_PATH = f"{ERROR}: [{error_prefix}] Invalid credentials. Credentials path can not be empty."
+ EMPTY_CREDENTIALS_STRING = f"{ERROR}: [{error_prefix}] Invalid credentials. Credentials string can not be empty."
+ EMPTY_TOKEN_VALUE = f"{ERROR}: [{error_prefix}] Invalid credentials. Token can not be empty."
+ EMPTY_API_KEY_VALUE = f"{ERROR}: [{error_prefix}] Invalid credentials. Api key can not be empty."
+ INVALID_API_KEY = f"{ERROR}: [{error_prefix}] Invalid credentials. Api key is invalid."
+
+ INVALID_BEARER_TOKEN = f"{ERROR}: [{error_prefix}] Bearer token is invalid or expired."
+ INVALID_CREDENTIALS_FILE = f"{ERROR}: [{error_prefix}] Credentials file is either null or an invalid file."
+ INVALID_CREDENTIALS_STRING_FORMAT = f"{ERROR}: [{error_prefix}] Credentials string in not in a valid JSON string format."
+ PRIVATE_KEY_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Private key is required."
+ CLIENT_ID_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Client ID is required."
+ KEY_ID_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Key ID is required."
+ TOKEN_URI_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Token URI is required."
+ INVALID_TOKEN_URI = f"{ERROR}: [{error_prefix}] Invalid value for token URI in credentials."
+ FAILED_TO_GET_BEARER_TOKEN = f"{ERROR}: [{error_prefix}] Failed to generate bearer token."
+ UNAUTHORIZED_ERROR_IN_GETTING_BEARER_TOKEN = f"{ERROR}: [{error_prefix}] Authorization failed while retrieving the bearer token."
+
+
+ TABLE_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Table is required."
+ EMPTY_TABLE_NAME =f"{ERROR}: [{error_prefix}] Invalid {{}} request. Table name can not be empty."
+ VALUES_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Values are required."
+ EMPTY_VALUES = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Values can not be empty."
+ EMPTY_OR_NULL_VALUE_IN_VALUES = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Value can not be null or empty in values for key {{}}."
+ EMPTY_OR_NULL_KEY_IN_VALUES = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Key can not be null or empty in values."
+ EMPTY_UPSERT = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Upsert can not be empty."
+ HOMOGENOUS_NOT_SUPPORTED_WITH_UPSERT = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Homogenous is not supported when upsert is passed."
+ EMPTY_TOKENS = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Tokens can not be empty."
+ EMPTY_OR_NULL_VALUE_IN_TOKENS = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Value can not be null or empty in tokens for key {{}}."
+ EMPTY_OR_NULL_KEY_IN_TOKENS = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Key can not be null or empty in tokens."
+ MISMATCH_OF_FIELDS_AND_TOKENS = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Keys for values and tokens are not matching."
+ FILE_UPLOAD_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] File upload failed."
+
+ EMPTY_IDS = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Ids can not be empty."
+ EMPTY_OR_NULL_ID_IN_IDS = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Id can not be null or empty in ids at index {{}}."
+ TOKENIZATION_NOT_SUPPORTED_WITH_REDACTION= f"{ERROR}: [{error_prefix}] Invalid {{}} request. Tokenization is not supported when redaction is applied."
+ TOKENIZATION_SUPPORTED_ONLY_WITH_IDS=f"{ERROR}: [{error_prefix}] Invalid {{}} request. Tokenization is not supported when column name and values are passed."
+ TOKENS_NOT_ALLOWED_WITH_BYOT_DISABLE = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Tokens are not allowed when token_mode is DISABLE."
+ INSUFFICIENT_TOKENS_PASSED_FOR_BYOT_ENABLE_STRICT =f"{ERROR}: [{error_prefix}] Invalid {{}} request. For token_mode as ENABLE_STRICT, tokens should be passed for all fields."
+ TOKENS_REQUIRED = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Tokens are required."
+ EMPTY_FIELDS = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Fields can not be empty."
+ EMPTY_OFFSET = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Offset ca not be empty."
+ NEITHER_IDS_NOR_COLUMN_NAME_PASSED = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Neither ids nor column name and values are passed."
+ BOTH_IDS_AND_COLUMN_NAME_PASSED = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Both ids and column name and values are passed."
+ COLUMN_NAME_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Column name is required when column values are passed."
+ COLUMN_VALUES_IS_REQUIRED_GET = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Column values are required when column name is passed."
+ SKYFLOW_ID_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Skyflow Id is required."
+ EMPTY_SKYFLOW_ID = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Skyflow Id can not be empty."
+
+ COLUMN_VALUES_IS_REQUIRED_TOKENIZE = f"{ERROR}: [{error_prefix}] Invalid {{}} request. column_values are required."
+ EMPTY_COLUMN_GROUP_IN_COLUMN_VALUES = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Column group can not be null or empty in column values at index %s2."
+
+ EMPTY_QUERY= f"{ERROR}: [{error_prefix}] Invalid {{}} request. Query can not be empty."
+ QUERY_IS_REQUIRED = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Query is required."
+
+ INSERT_RECORDS_REJECTED = f"{ERROR}: [{error_prefix}] Insert call resulted in failure."
+ DETOKENIZE_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Detokenize request resulted in failure."
+ DELETE_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Delete request resulted in failure."
+ TOKENIZE_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Tokenize request resulted in failure."
+ UPDATE_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Update request resulted in failure."
+ QUERY_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Query request resulted in failure."
+ GET_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Get request resulted in failure."
+ INVOKE_CONNECTION_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Invoke connection request resulted in failure."
+
+ EMPTY_RUN_ID = f"{ERROR}: [{error_prefix}] Validation error. Run id cannot be empty. Specify a valid run id."
+ INVALID_RUN_ID = f"{ERROR}: [{error_prefix}] Validation error. Invalid run id. Specify a valid run id as string."
+ DEIDENTIFY_FILE_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Deidentify file resulted in failure."
+ DETECT_RUN_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Detect get run resulted in failure."
+ DEIDENTIFY_TEXT_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Deidentify text resulted in failure."
+ SAVING_DEIDENTIFY_FILE_FAILED = f"{ERROR}: [{error_prefix}] Error while saving deidentified file to output directory."
+ REIDENTIFY_TEXT_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Reidentify text resulted in failure."
+ DETECT_FILE_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Deidentify file resulted in failure."
+ EMPTY_FILE_COLUMN_NAME = f"{ERROR}: [{error_prefix}] Empty column name in FILE_UPLOAD"
+
+ class Interface(Enum):
+ INSERT = "INSERT"
+ GET = "GET"
+ QUERY = "QUERY"
+ DETOKENIZE = " DETOKENIZE"
+ TOKENIZE = "TOKENIZE"
+ UPDATE = "UPDATE"
+ DELETE = "DELETE"
+
+ class HttpStatus(Enum):
+ BAD_REQUEST = "Bad Request"
+
+ class Warning(Enum):
+ DETOKENIZE_REDACTION_KEY_DEPRECATED = (
+ f"{WARN}: [{error_prefix}] 'redaction' key in detokenize data is deprecated and will be removed in a future version. Use 'redaction_type' instead."
+ )
+ UPDATE_LOG_LEVEL_DEPRECATED = (
+ f"{WARN}: [{error_prefix}] Skyflow.update_log_level() is deprecated. "
+ "Use Skyflow.set_log_level() instead."
+ )
+ FILE_UPLOAD_REQUEST_ARG_ORDER_DEPRECATED = (
+ f"{WARN}: [{error_prefix}] FileUploadRequest: argument order changed. "
+ "Old positional order: (table, skyflow_id, column_name). "
+ "New order: FileUploadRequest(table, column_name=..., skyflow_id=...)."
+ )
+
+
+
diff --git a/common/utils/_utils.py b/common/utils/_utils.py
new file mode 100644
index 00000000..4a1c0637
--- /dev/null
+++ b/common/utils/_utils.py
@@ -0,0 +1,50 @@
+import os
+import re
+
+import dotenv
+from dotenv import load_dotenv
+
+from common.errors import SkyflowError
+from . import SkyflowMessages
+from .constants import PROTOCOL, ApiKey
+from .enums import Env, EnvUrls
+from .logger import log_error_log
+
+invalid_input_error_code = SkyflowMessages.ErrorCodes.INVALID_INPUT.value
+
+
+def get_credentials(config_level_creds=None, common_skyflow_creds=None, logger=None):
+ if config_level_creds is not None:
+ return config_level_creds
+ if common_skyflow_creds is not None:
+ return common_skyflow_creds
+ dotenv_path = dotenv.find_dotenv(usecwd=True)
+ if dotenv_path:
+ load_dotenv(dotenv_path)
+ env_skyflow_credentials = os.getenv("SKYFLOW_CREDENTIALS")
+ if env_skyflow_credentials:
+ env_creds = env_skyflow_credentials.strip().replace('\n', '\\n')
+ return {'credentials_string': env_creds}
+ raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIALS.value, invalid_input_error_code)
+
+
+def validate_api_key(api_key: str, logger=None) -> bool:
+ if len(api_key) != ApiKey.LENGTH:
+ log_error_log(SkyflowMessages.ErrorLogs.INVALID_API_KEY.value, logger=logger)
+ return False
+ api_key_pattern = re.compile(r'^sky-[a-zA-Z0-9]{5}-[a-fA-F0-9]{32}$')
+
+ return bool(api_key_pattern.match(api_key))
+
+
+def get_vault_url(cluster_id, env, vault_id, logger=None):
+ if not cluster_id or not isinstance(cluster_id, str) or not cluster_id.strip():
+ raise SkyflowError(SkyflowMessages.Error.INVALID_CLUSTER_ID.value.format(vault_id), invalid_input_error_code)
+
+ if env not in Env:
+ raise SkyflowError(SkyflowMessages.Error.INVALID_ENV.value.format(vault_id), invalid_input_error_code)
+
+ base_url = EnvUrls[env.name].value
+ protocol = PROTOCOL
+
+ return f"{protocol}://{cluster_id}.{base_url}"
diff --git a/skyflow/utils/constants.py b/common/utils/constants.py
similarity index 100%
rename from skyflow/utils/constants.py
rename to common/utils/constants.py
diff --git a/skyflow/utils/enums/__init__.py b/common/utils/enums/__init__.py
similarity index 100%
rename from skyflow/utils/enums/__init__.py
rename to common/utils/enums/__init__.py
diff --git a/skyflow/utils/enums/content_types.py b/common/utils/enums/content_types.py
similarity index 100%
rename from skyflow/utils/enums/content_types.py
rename to common/utils/enums/content_types.py
diff --git a/skyflow/utils/enums/detect_entities.py b/common/utils/enums/detect_entities.py
similarity index 100%
rename from skyflow/utils/enums/detect_entities.py
rename to common/utils/enums/detect_entities.py
diff --git a/skyflow/utils/enums/detect_output_transcriptions.py b/common/utils/enums/detect_output_transcriptions.py
similarity index 100%
rename from skyflow/utils/enums/detect_output_transcriptions.py
rename to common/utils/enums/detect_output_transcriptions.py
diff --git a/skyflow/utils/enums/env.py b/common/utils/enums/env.py
similarity index 100%
rename from skyflow/utils/enums/env.py
rename to common/utils/enums/env.py
diff --git a/skyflow/utils/enums/log_level.py b/common/utils/enums/log_level.py
similarity index 100%
rename from skyflow/utils/enums/log_level.py
rename to common/utils/enums/log_level.py
diff --git a/skyflow/utils/enums/masking_method.py b/common/utils/enums/masking_method.py
similarity index 100%
rename from skyflow/utils/enums/masking_method.py
rename to common/utils/enums/masking_method.py
diff --git a/skyflow/utils/enums/redaction_type.py b/common/utils/enums/redaction_type.py
similarity index 100%
rename from skyflow/utils/enums/redaction_type.py
rename to common/utils/enums/redaction_type.py
diff --git a/skyflow/utils/enums/request_method.py b/common/utils/enums/request_method.py
similarity index 100%
rename from skyflow/utils/enums/request_method.py
rename to common/utils/enums/request_method.py
diff --git a/skyflow/utils/enums/token_mode.py b/common/utils/enums/token_mode.py
similarity index 100%
rename from skyflow/utils/enums/token_mode.py
rename to common/utils/enums/token_mode.py
diff --git a/skyflow/utils/enums/token_type.py b/common/utils/enums/token_type.py
similarity index 100%
rename from skyflow/utils/enums/token_type.py
rename to common/utils/enums/token_type.py
diff --git a/skyflow/utils/logger/__init__.py b/common/utils/logger/__init__.py
similarity index 100%
rename from skyflow/utils/logger/__init__.py
rename to common/utils/logger/__init__.py
diff --git a/skyflow/utils/logger/_log_helpers.py b/common/utils/logger/_log_helpers.py
similarity index 100%
rename from skyflow/utils/logger/_log_helpers.py
rename to common/utils/logger/_log_helpers.py
diff --git a/skyflow/utils/logger/_logger.py b/common/utils/logger/_logger.py
similarity index 100%
rename from skyflow/utils/logger/_logger.py
rename to common/utils/logger/_logger.py
diff --git a/common/utils/validations/__init__.py b/common/utils/validations/__init__.py
new file mode 100644
index 00000000..d49cc5de
--- /dev/null
+++ b/common/utils/validations/__init__.py
@@ -0,0 +1,9 @@
+from ._validations import (
+ validate_required_field,
+ validate_api_key,
+ validate_credentials,
+ validate_log_level,
+ validate_keys,
+ validate_vault_config,
+ validate_update_vault_config,
+)
diff --git a/common/utils/validations/_validations.py b/common/utils/validations/_validations.py
new file mode 100644
index 00000000..f600b6a0
--- /dev/null
+++ b/common/utils/validations/_validations.py
@@ -0,0 +1,233 @@
+from common.errors import SkyflowError
+from common.service_account import is_expired
+from common.utils import SkyflowMessages
+from common.utils.constants import ApiKey, ConfigField, CredentialField, OptionField
+from common.utils.enums import Env, LogLevel
+from common.utils.logger import log_error_log, log_info
+from common.utils._helpers import is_valid_url
+
+invalid_input_error_code = SkyflowMessages.ErrorCodes.INVALID_INPUT.value
+
+VALID_VAULT_CONFIG_KEYS = [
+ ConfigField.VAULT_ID,
+ ConfigField.CLUSTER_ID,
+ ConfigField.CREDENTIALS,
+ ConfigField.ENV,
+]
+
+
+def validate_required_field(logger, config, field_name, expected_type, empty_error, invalid_error, messages=None):
+ messages = messages or SkyflowMessages
+ field_value = config.get(field_name)
+
+ if field_name not in config or not isinstance(field_value, expected_type):
+ if field_name == ConfigField.VAULT_ID:
+ log_error_log(messages.ErrorLogs.VAULTID_IS_REQUIRED.value, logger)
+ if field_name == ConfigField.CLUSTER_ID:
+ log_error_log(messages.ErrorLogs.CLUSTER_ID_IS_REQUIRED.value, logger)
+ if field_name == OptionField.CONNECTION_ID:
+ log_error_log(messages.ErrorLogs.CONNECTION_ID_IS_REQUIRED.value, logger)
+ if field_name == OptionField.CONNECTION_URL:
+ log_error_log(messages.ErrorLogs.INVALID_CONNECTION_URL.value, logger)
+ raise SkyflowError(invalid_error, invalid_input_error_code)
+
+ if isinstance(field_value, str) and not field_value.strip():
+ if field_name == ConfigField.VAULT_ID:
+ log_error_log(messages.ErrorLogs.EMPTY_VAULTID.value, logger)
+ if field_name == ConfigField.CLUSTER_ID:
+ log_error_log(messages.ErrorLogs.EMPTY_CLUSTER_ID.value, logger)
+ if field_name == OptionField.CONNECTION_ID:
+ log_error_log(messages.ErrorLogs.EMPTY_CONNECTION_ID.value, logger)
+ if field_name == OptionField.CONNECTION_URL:
+ log_error_log(messages.ErrorLogs.EMPTY_CONNECTION_URL.value, logger)
+ if field_name == CredentialField.PATH:
+ log_error_log(messages.ErrorLogs.EMPTY_CREDENTIALS_PATH.value, logger)
+ if field_name == CredentialField.CREDENTIALS_STRING:
+ log_error_log(messages.ErrorLogs.EMPTY_CREDENTIALS_STRING.value, logger)
+ if field_name == CredentialField.TOKEN:
+ log_error_log(messages.ErrorLogs.EMPTY_TOKEN_VALUE.value, logger)
+ if field_name == CredentialField.API_KEY:
+ log_error_log(messages.ErrorLogs.EMPTY_API_KEY_VALUE.value, logger)
+ raise SkyflowError(empty_error, invalid_input_error_code)
+
+
+def validate_api_key(api_key: str, logger=None, messages=None) -> bool:
+ messages = messages or SkyflowMessages
+ if not api_key.startswith(ApiKey.SKY_PREFIX):
+ log_error_log(messages.ErrorLogs.INVALID_API_KEY.value, logger=logger)
+ return False
+
+ if len(api_key) != ApiKey.LENGTH:
+ log_error_log(messages.ErrorLogs.INVALID_API_KEY.value, logger=logger)
+ return False
+
+ return True
+
+
+def validate_credentials(logger, credentials, config_id_type=None, config_id=None, messages=None):
+ messages = messages or SkyflowMessages
+ key_present = [k for k in [CredentialField.PATH, CredentialField.TOKEN, CredentialField.CREDENTIALS_STRING, CredentialField.API_KEY] if credentials.get(k)]
+
+ if len(key_present) == 0:
+ error_message = (
+ messages.Error.INVALID_CREDENTIALS_IN_CONFIG.value.format(config_id_type, config_id)
+ if config_id_type and config_id else
+ messages.Error.INVALID_CREDENTIALS.value
+ )
+ log_error_log(error_message, logger)
+ raise SkyflowError(error_message, invalid_input_error_code)
+ elif len(key_present) > 1:
+ error_message = (
+ messages.Error.MULTIPLE_CREDENTIALS_PASSED_IN_CONFIG.value.format(config_id_type, config_id)
+ if config_id_type and config_id else
+ messages.Error.MULTIPLE_CREDENTIALS_PASSED.value
+ )
+ log_error_log(error_message, logger)
+ raise SkyflowError(error_message, invalid_input_error_code)
+
+ if CredentialField.ROLES in credentials:
+ validate_required_field(
+ logger, credentials, CredentialField.ROLES, list,
+ messages.Error.INVALID_ROLES_KEY_TYPE_IN_CONFIG.value.format(config_id_type, config_id)
+ if config_id_type and config_id else messages.Error.INVALID_ROLES_KEY_TYPE.value,
+ messages.Error.EMPTY_ROLES_IN_CONFIG.value.format(config_id_type, config_id)
+ if config_id_type and config_id else messages.Error.EMPTY_ROLES.value,
+ messages=messages)
+
+ if CredentialField.CONTEXT in credentials:
+ validate_required_field(
+ logger, credentials, CredentialField.CONTEXT, str,
+ messages.Error.EMPTY_CONTEXT_IN_CONFIG.value.format(config_id_type, config_id)
+ if config_id_type and config_id else messages.Error.EMPTY_CONTEXT.value,
+ messages.Error.INVALID_CONTEXT_IN_CONFIG.value.format(config_id_type, config_id)
+ if config_id_type and config_id else messages.Error.INVALID_CONTEXT.value,
+ messages=messages)
+
+ if CredentialField.CREDENTIALS_STRING in credentials:
+ validate_required_field(
+ logger, credentials, CredentialField.CREDENTIALS_STRING, str,
+ messages.Error.EMPTY_CREDENTIALS_STRING_IN_CONFIG.value.format(config_id_type, config_id)
+ if config_id_type and config_id else messages.Error.EMPTY_CREDENTIALS_STRING.value,
+ messages.Error.INVALID_CREDENTIALS_STRING_IN_CONFIG.value.format(config_id_type, config_id)
+ if config_id_type and config_id else messages.Error.INVALID_CREDENTIALS_STRING.value,
+ messages=messages)
+ elif CredentialField.PATH in credentials:
+ validate_required_field(
+ logger, credentials, CredentialField.PATH, str,
+ messages.Error.EMPTY_CREDENTIAL_FILE_PATH_IN_CONFIG.value.format(config_id_type, config_id)
+ if config_id_type and config_id else messages.Error.EMPTY_CREDENTIAL_FILE_PATH.value,
+ messages.Error.INVALID_CREDENTIAL_FILE_PATH_IN_CONFIG.value.format(config_id_type, config_id)
+ if config_id_type and config_id else messages.Error.INVALID_CREDENTIAL_FILE_PATH.value,
+ messages=messages)
+ elif CredentialField.TOKEN in credentials:
+ validate_required_field(
+ logger, credentials, CredentialField.TOKEN, str,
+ messages.Error.EMPTY_CREDENTIALS_TOKEN.value.format(config_id_type, config_id)
+ if config_id_type and config_id else messages.Error.EMPTY_CREDENTIALS_TOKEN.value,
+ messages.Error.INVALID_CREDENTIALS_TOKEN.value.format(config_id_type, config_id)
+ if config_id_type and config_id else messages.Error.INVALID_CREDENTIALS_TOKEN.value,
+ messages=messages)
+ if is_expired(credentials.get(CredentialField.TOKEN), logger):
+ log_error_log(messages.ErrorLogs.INVALID_BEARER_TOKEN.value, logger)
+ raise SkyflowError(
+ messages.Error.EXPIRED_BEARER_TOKEN.value
+ if config_id_type and config_id else messages.Error.EXPIRED_BEARER_TOKEN.value,
+ invalid_input_error_code
+ )
+ elif CredentialField.API_KEY in credentials:
+ validate_required_field(
+ logger, credentials, CredentialField.API_KEY, str,
+ messages.Error.EMPTY_API_KEY.value.format(config_id_type, config_id)
+ if config_id_type and config_id else messages.Error.EMPTY_API_KEY.value,
+ messages.Error.INVALID_API_KEY.value.format(config_id_type, config_id)
+ if config_id_type and config_id else messages.Error.INVALID_API_KEY.value,
+ messages=messages)
+ if not validate_api_key(credentials.get(CredentialField.API_KEY), logger, messages=messages):
+ raise SkyflowError(messages.Error.INVALID_API_KEY.value.format(config_id_type, config_id)
+ if config_id_type and config_id else messages.Error.INVALID_API_KEY.value,
+ invalid_input_error_code)
+
+ if CredentialField.TOKEN_URI_OPTION in credentials:
+ token_uri = credentials.get(CredentialField.TOKEN_URI_OPTION)
+ if (
+ token_uri is None
+ or not isinstance(token_uri, str)
+ or not is_valid_url(token_uri)
+ ):
+ log_error_log(messages.ErrorLogs.INVALID_TOKEN_URI.value, logger)
+ raise SkyflowError(messages.Error.INVALID_TOKEN_URI.value, invalid_input_error_code)
+
+
+def validate_log_level(logger, log_level, messages=None):
+ messages = messages or SkyflowMessages
+ if not isinstance(log_level, LogLevel):
+ log_error_log(messages.ErrorLogs.INVALID_LOG_LEVEL.value, logger)
+ raise SkyflowError(messages.Error.INVALID_LOG_LEVEL.value, invalid_input_error_code)
+
+
+def validate_keys(logger, config, config_keys, messages=None):
+ messages = messages or SkyflowMessages
+ for key in config.keys():
+ if key not in config_keys:
+ log_error_log(messages.ErrorLogs.INVALID_KEY.value.format(key), logger)
+ raise SkyflowError(messages.Error.INVALID_KEY.value.format(key), invalid_input_error_code)
+
+
+def validate_vault_config(logger, config, messages=None):
+ messages = messages or SkyflowMessages
+ log_info(messages.Info.VALIDATING_VAULT_CONFIG.value, logger)
+ validate_keys(logger, config, VALID_VAULT_CONFIG_KEYS, messages=messages)
+
+ validate_required_field(
+ logger, config, ConfigField.VAULT_ID, str,
+ messages.Error.EMPTY_VAULT_ID.value,
+ messages.Error.INVALID_VAULT_ID.value,
+ messages=messages
+ )
+ vault_id = config.get(ConfigField.VAULT_ID)
+
+ validate_required_field(
+ logger, config, ConfigField.CLUSTER_ID, str,
+ messages.Error.EMPTY_CLUSTER_ID.value.format(vault_id),
+ messages.Error.INVALID_CLUSTER_ID.value.format(vault_id),
+ messages=messages
+ )
+
+ if ConfigField.CREDENTIALS in config and not config.get(ConfigField.CREDENTIALS):
+ raise SkyflowError(messages.Error.EMPTY_CREDENTIALS.value.format("vault", vault_id), invalid_input_error_code)
+
+ if ConfigField.CREDENTIALS in config and config.get(ConfigField.CREDENTIALS):
+ validate_credentials(logger, config.get(ConfigField.CREDENTIALS), "vault", vault_id, messages=messages)
+
+ if ConfigField.ENV in config and config.get(ConfigField.ENV) not in Env:
+ log_error_log(messages.ErrorLogs.ENV_IS_REQUIRED.value, logger)
+ raise SkyflowError(messages.Error.INVALID_ENV.value.format(vault_id), invalid_input_error_code)
+
+ return True
+
+
+def validate_update_vault_config(logger, config, messages=None):
+ """Credentials are required on update (unlike on initial add, where they're optional)."""
+ messages = messages or SkyflowMessages
+ validate_keys(logger, config, VALID_VAULT_CONFIG_KEYS, messages=messages)
+
+ validate_required_field(
+ logger, config, ConfigField.VAULT_ID, str,
+ messages.Error.EMPTY_VAULT_ID.value,
+ messages.Error.INVALID_VAULT_ID.value,
+ messages=messages
+ )
+ vault_id = config.get(ConfigField.VAULT_ID)
+
+ if ConfigField.CLUSTER_ID in config and not config.get(ConfigField.CLUSTER_ID):
+ raise SkyflowError(messages.Error.INVALID_CLUSTER_ID.value.format(vault_id), invalid_input_error_code)
+
+ if ConfigField.ENV in config and config.get(ConfigField.ENV) not in Env:
+ raise SkyflowError(messages.Error.INVALID_ENV.value.format(vault_id), invalid_input_error_code)
+
+ if ConfigField.CREDENTIALS not in config:
+ raise SkyflowError(messages.Error.EMPTY_CREDENTIALS.value.format("vault", vault_id), invalid_input_error_code)
+
+ validate_credentials(logger, config.get(ConfigField.CREDENTIALS), "vault", vault_id, messages=messages)
+
+ return True
diff --git a/common/vault/base_vault_client.py b/common/vault/base_vault_client.py
new file mode 100644
index 00000000..8ba044ac
--- /dev/null
+++ b/common/vault/base_vault_client.py
@@ -0,0 +1,116 @@
+from abc import ABC, abstractmethod
+
+from common.service_account import generate_bearer_token, generate_bearer_token_from_creds, is_expired
+from common.utils import get_credentials, SkyflowMessages
+from common.utils.logger import log_info
+from common.utils.constants import OptionField, CredentialField, ConfigField
+
+
+class IVaultClient(ABC):
+ @abstractmethod
+ def resolve_vault_url(self, cluster_id, env, vault_id, logger=None):
+ raise NotImplementedError
+
+ @abstractmethod
+ def initialize_api_client(self, vault_url, bearer_token):
+ raise NotImplementedError
+
+
+class BaseVaultClient(IVaultClient):
+ def __init__(self, config):
+ self._config = config
+ self._common_skyflow_credentials = None
+ self._log_level = None
+ self._api_client = None
+ self._logger = None
+ self._is_config_updated = False
+ self._bearer_token = None
+ self._credentials = None
+ self._vault_url = None
+ self._is_static_token = None
+
+ def set_common_skyflow_credentials(self, credentials):
+ self._common_skyflow_credentials = credentials
+
+ def set_logger(self, log_level, logger):
+ self._log_level = log_level
+ self._logger = logger
+
+ def initialize_client_configuration(self):
+ if self._api_client is not None and not self._is_config_updated:
+ if self._is_static_token:
+ return
+ if self._bearer_token is not None and not is_expired(self._bearer_token):
+ return
+
+ needs_reinit = self._api_client is None or self._is_config_updated
+ if needs_reinit:
+ self._credentials = get_credentials(self._config.get(ConfigField.CREDENTIALS), self._common_skyflow_credentials, logger=self._logger)
+ self._vault_url = self.resolve_vault_url(self._config.get(ConfigField.CLUSTER_ID),
+ self._config.get(ConfigField.ENV),
+ self._config.get(ConfigField.VAULT_ID),
+ logger=self._logger)
+ self._is_static_token = CredentialField.TOKEN in self._credentials or CredentialField.API_KEY in self._credentials
+ bearer_token = self.get_bearer_token(self._credentials)
+ self._bearer_token = bearer_token
+ if needs_reinit:
+ self.initialize_api_client(self._vault_url, bearer_token)
+
+ def get_current_bearer_token(self):
+ return self._bearer_token
+
+ def get_current_vault_url(self):
+ return self._vault_url
+
+ def get_vault_id(self):
+ return self._config.get(ConfigField.VAULT_ID)
+
+ def get_bearer_token(self, credentials):
+ if CredentialField.API_KEY in credentials:
+ return credentials.get(CredentialField.API_KEY)
+ elif CredentialField.TOKEN in credentials:
+ return credentials.get(CredentialField.TOKEN)
+
+ options = {
+ OptionField.ROLE_IDS: self._config.get(OptionField.ROLES),
+ OptionField.CTX: self._config.get(OptionField.CTX)
+ }
+ if CredentialField.TOKEN_URI_OPTION in credentials and credentials.get(CredentialField.TOKEN_URI_OPTION):
+ options[CredentialField.TOKEN_URI_OPTION] = credentials.get(CredentialField.TOKEN_URI_OPTION)
+
+ if self._bearer_token is None or self._is_config_updated or is_expired(self._bearer_token):
+ if CredentialField.PATH in credentials:
+ self._bearer_token, _ = generate_bearer_token(
+ credentials.get(CredentialField.PATH),
+ options,
+ self._logger
+ )
+ else:
+ credentials_string = credentials.get(CredentialField.CREDENTIALS_STRING)
+ log_info(SkyflowMessages.Info.GENERATE_BEARER_TOKEN_FROM_CREDENTIALS_STRING_TRIGGERED.value, self._logger)
+ self._bearer_token, _ = generate_bearer_token_from_creds(
+ credentials_string,
+ options,
+ self._logger
+ )
+ self._is_config_updated = False
+ else:
+ log_info(SkyflowMessages.Info.REUSE_BEARER_TOKEN.value, self._logger)
+
+ return self._bearer_token
+
+ def update_config(self, config):
+ self._config.update(config)
+ self._is_config_updated = True
+
+ def get_config(self):
+ return self._config
+
+ def get_common_skyflow_credentials(self):
+ return self._common_skyflow_credentials
+
+ def get_log_level(self):
+ return self._log_level
+
+ def get_logger(self):
+ return self._logger
diff --git a/common/vault/base_vault_controller.py b/common/vault/base_vault_controller.py
new file mode 100644
index 00000000..f95dce11
--- /dev/null
+++ b/common/vault/base_vault_controller.py
@@ -0,0 +1,62 @@
+from abc import ABC, abstractmethod
+
+from common.errors import SkyflowError
+from common.utils import SkyflowMessages as _CommonSkyflowMessages
+from common.vault.data import BaseInsertRequest, BaseInsertResponse
+
+_INVALID_INPUT_ERROR_CODE = _CommonSkyflowMessages.ErrorCodes.INVALID_INPUT.value
+
+
+class IVaultController(ABC):
+
+ @abstractmethod
+ def insert(self, request: BaseInsertRequest) -> BaseInsertResponse:
+ raise NotImplementedError
+
+ @abstractmethod
+ def get(self, request):
+ raise NotImplementedError
+
+ @abstractmethod
+ def update(self, request):
+ raise NotImplementedError
+
+ @abstractmethod
+ def delete(self, request):
+ raise NotImplementedError
+
+ @abstractmethod
+ def query(self, request):
+ raise NotImplementedError
+
+ @abstractmethod
+ def detokenize(self, request):
+ raise NotImplementedError
+
+
+class BaseVaultController(IVaultController):
+
+ _skyflow_messages = None
+
+ def __init__(self, vault_client):
+ self._vault_client = vault_client
+
+ def _validate_table_name_if_present(self, table):
+ if table is not None and (not isinstance(table, str) or not table.strip()):
+ raise SkyflowError(
+ self._skyflow_messages.Error.INVALID_TABLE_NAME_IN_INSERT.value,
+ _INVALID_INPUT_ERROR_CODE,
+ )
+
+ def _validate_field_values(self, values):
+ if not isinstance(values, dict) or not values:
+ raise SkyflowError(
+ self._skyflow_messages.Error.INVALID_RECORD_DATA_IN_INSERT.value,
+ _INVALID_INPUT_ERROR_CODE,
+ )
+ for key, value in values.items():
+ if not isinstance(key, str) or not key.strip():
+ raise SkyflowError(
+ self._skyflow_messages.Error.EMPTY_KEY_IN_INSERT_DATA.value,
+ _INVALID_INPUT_ERROR_CODE,
+ )
diff --git a/common/vault/data/__init__.py b/common/vault/data/__init__.py
new file mode 100644
index 00000000..23b585aa
--- /dev/null
+++ b/common/vault/data/__init__.py
@@ -0,0 +1,2 @@
+from ._base_insert_request import BaseInsertRequest
+from ._base_insert_response import BaseInsertResponse
diff --git a/common/vault/data/_base_insert_request.py b/common/vault/data/_base_insert_request.py
new file mode 100644
index 00000000..4172fd69
--- /dev/null
+++ b/common/vault/data/_base_insert_request.py
@@ -0,0 +1,7 @@
+from typing import Union
+
+class BaseInsertRequest:
+ def __init__(self, table: str, values: list, upsert: Union[str, dict] = None):
+ self.table = table
+ self.values = values
+ self.upsert = upsert
\ No newline at end of file
diff --git a/common/vault/data/_base_insert_response.py b/common/vault/data/_base_insert_response.py
new file mode 100644
index 00000000..b0b60377
--- /dev/null
+++ b/common/vault/data/_base_insert_response.py
@@ -0,0 +1,11 @@
+class BaseInsertResponse:
+
+ def __init__(self, inserted_fields=None, errors=None):
+ self.inserted_fields = inserted_fields
+ self.errors = errors
+
+ def __repr__(self):
+ return f"{type(self).__name__}(inserted_fields={self.inserted_fields}, errors={self.errors})"
+
+ def __str__(self):
+ return self.__repr__()
diff --git a/flowvault/requirements.txt b/flowvault/requirements.txt
new file mode 100644
index 00000000..7d2fb3a8
--- /dev/null
+++ b/flowvault/requirements.txt
@@ -0,0 +1,5 @@
+httpx>=0.21.2
+pydantic>= 1.9.2
+pydantic-core>=2.18.2
+typing_extensions>= 4.0.0
+coverage >= 7.8.0
diff --git a/flowvault/setup.py b/flowvault/setup.py
new file mode 100644
index 00000000..e3fa571e
--- /dev/null
+++ b/flowvault/setup.py
@@ -0,0 +1,84 @@
+'''
+ Copyright (c) 2022 Skyflow, Inc.
+'''
+import os
+import shutil
+import sys
+
+from setuptools import setup, find_packages
+from setuptools.command.build_py import build_py as _build_py
+
+
+if sys.version_info < (3, 9):
+ raise RuntimeError("skyflow requires Python 3.9+")
+current_version = '1.0.0'
+
+HERE = os.path.abspath(os.path.dirname(__file__))
+REPO_ROOT = os.path.dirname(HERE)
+COMMON_SRC = os.path.join(REPO_ROOT, 'common')
+
+with open(os.path.join(REPO_ROOT, 'README.md'), 'r', encoding='utf-8') as f:
+ long_description = f.read()
+
+_COMMON_EXCLUDE_DIRS = {'__pycache__', '.pytest_cache', 'tests', '.mypy_cache'}
+_COMMON_EXCLUDE_FILES = {'setup.py', 'pyproject.toml', 'requirements.txt', '.gitignore'}
+_COMMON_EXCLUDE_SUFFIXES = ('.egg-info',)
+
+
+def _ignore_common_files(_directory, names):
+ ignored = set()
+ for name in names:
+ if name in _COMMON_EXCLUDE_DIRS or name in _COMMON_EXCLUDE_FILES:
+ ignored.add(name)
+ elif name.endswith(_COMMON_EXCLUDE_SUFFIXES):
+ ignored.add(name)
+ return ignored
+
+
+class CustomBuildPy(_build_py):
+ """SK-2938 Option C bundling mechanism -- see v2/setup.py for full rationale. Bundles the
+ sibling common/ source tree into this variant's wheel; wheel builds only, not sdist."""
+
+ def run(self):
+ super().run()
+ dest = os.path.join(self.build_lib, 'common')
+ if os.path.exists(dest):
+ shutil.rmtree(dest)
+ shutil.copytree(COMMON_SRC, dest, ignore=_ignore_common_files)
+
+
+setup(
+ name='skyflow-flowvault',
+ version=current_version,
+ author='Skyflow',
+ author_email='service-ops@skyflow.com',
+ packages=find_packages(where='.', exclude=['test*', 'samples*']),
+ package_data={
+ 'skyflow_flowvault': ['py.typed'],
+ 'skyflow_flowvault.generated.rest': ['py.typed'],
+ },
+ cmdclass={'build_py': CustomBuildPy},
+ url='https://github.com/skyflowapi/skyflow-python/',
+ license='LICENSE',
+ description='Skyflow SDK for the Python programming language (v3 / flowservice API)',
+ long_description=long_description,
+ long_description_content_type='text/markdown',
+ install_requires=[
+ 'pydantic >= 2.0.0',
+ 'typing-extensions >= 4.0.0',
+ 'PyJWT >= 2.12, < 3',
+ 'cryptography >= 44.0.2',
+ 'httpx >= 0.21.2',
+ 'python-dotenv >= 1.1.0, < 2',
+ # NOTE: 'requests' intentionally omitted -- only used today by v2's Connection
+ # controller, which isn't part of v3's scope this round.
+ ],
+ extras_require={
+ 'dev': [
+ 'codespell >= 2.4.1',
+ 'ruff >= 0.9.0',
+ 'pre-commit >= 4.3.0',
+ ]
+ },
+ python_requires=">=3.9",
+)
diff --git a/skyflow/__init__.py b/flowvault/skyflow_flowvault/__init__.py
similarity index 100%
rename from skyflow/__init__.py
rename to flowvault/skyflow_flowvault/__init__.py
diff --git a/skyflow/client/__init__.py b/flowvault/skyflow_flowvault/client/__init__.py
similarity index 100%
rename from skyflow/client/__init__.py
rename to flowvault/skyflow_flowvault/client/__init__.py
diff --git a/flowvault/skyflow_flowvault/client/skyflow.py b/flowvault/skyflow_flowvault/client/skyflow.py
new file mode 100644
index 00000000..59ff4d5f
--- /dev/null
+++ b/flowvault/skyflow_flowvault/client/skyflow.py
@@ -0,0 +1,10 @@
+from common.client.utils import make_skyflow_class
+from common.utils import SkyflowMessages
+from skyflow_flowvault.vault.client.client import VaultClient
+from skyflow_flowvault.vault.controller import VaultController
+
+Skyflow = make_skyflow_class(
+ vault_client_cls=VaultClient,
+ vault_controller_cls=VaultController,
+ skyflow_messages=SkyflowMessages,
+)
diff --git a/flowvault/skyflow_flowvault/error/__init__.py b/flowvault/skyflow_flowvault/error/__init__.py
new file mode 100644
index 00000000..dd3ea070
--- /dev/null
+++ b/flowvault/skyflow_flowvault/error/__init__.py
@@ -0,0 +1,3 @@
+from common.errors import SkyflowError
+
+__all__ = ["SkyflowError"]
diff --git a/tests/vault/__init__.py b/flowvault/skyflow_flowvault/generated/__init__.py
similarity index 100%
rename from tests/vault/__init__.py
rename to flowvault/skyflow_flowvault/generated/__init__.py
diff --git a/flowvault/skyflow_flowvault/generated/rest/__init__.py b/flowvault/skyflow_flowvault/generated/rest/__init__.py
new file mode 100644
index 00000000..14ae395a
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/__init__.py
@@ -0,0 +1,77 @@
+# This file was auto-generated by Fern from our API Definition.
+
+# isort: skip_file
+
+from .types import (
+ FlowEnumUpdateType,
+ FlowTokenizeResponseObjectToken,
+ GoogleprotobufAny,
+ ProtobufNullValue,
+ RpcStatus,
+ V1ColumnRedactions,
+ V1DeleteResponse,
+ V1DeleteResponseObject,
+ V1DeleteTokenResponseObject,
+ V1ExecuteQueryRecordResponse,
+ V1ExecuteQueryResponse,
+ V1ExecuteQueryResponseMetadata,
+ V1FlowDeleteTokenResponse,
+ V1FlowDetokenizeResponse,
+ V1FlowDetokenizeResponseObject,
+ V1FlowTokenizeRequestObject,
+ V1FlowTokenizeResponse,
+ V1FlowTokenizeResponseObject,
+ V1FlowVaultMetricsData,
+ V1FlowVaultMetricsResponse,
+ V1GetRequestData,
+ V1GetResponse,
+ V1InsertRecordData,
+ V1InsertResponse,
+ V1RecordResponseObject,
+ V1TokenGroupRedactions,
+ V1UniqueValue,
+ V1UpdateRecordData,
+ V1UpdateResponse,
+ V1Upsert,
+)
+from . import flowservice, records
+from .client import AsyncSkyflowAuth, SkyflowAuth
+from .version import __version__
+
+__all__ = [
+ "AsyncSkyflowAuth",
+ "FlowEnumUpdateType",
+ "FlowTokenizeResponseObjectToken",
+ "GoogleprotobufAny",
+ "ProtobufNullValue",
+ "RpcStatus",
+ "SkyflowAuth",
+ "V1ColumnRedactions",
+ "V1DeleteResponse",
+ "V1DeleteResponseObject",
+ "V1DeleteTokenResponseObject",
+ "V1ExecuteQueryRecordResponse",
+ "V1ExecuteQueryResponse",
+ "V1ExecuteQueryResponseMetadata",
+ "V1FlowDeleteTokenResponse",
+ "V1FlowDetokenizeResponse",
+ "V1FlowDetokenizeResponseObject",
+ "V1FlowTokenizeRequestObject",
+ "V1FlowTokenizeResponse",
+ "V1FlowTokenizeResponseObject",
+ "V1FlowVaultMetricsData",
+ "V1FlowVaultMetricsResponse",
+ "V1GetRequestData",
+ "V1GetResponse",
+ "V1InsertRecordData",
+ "V1InsertResponse",
+ "V1RecordResponseObject",
+ "V1TokenGroupRedactions",
+ "V1UniqueValue",
+ "V1UpdateRecordData",
+ "V1UpdateResponse",
+ "V1Upsert",
+ "__version__",
+ "flowservice",
+ "records",
+]
diff --git a/flowvault/skyflow_flowvault/generated/rest/client.py b/flowvault/skyflow_flowvault/generated/rest/client.py
new file mode 100644
index 00000000..0dc7a48c
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/client.py
@@ -0,0 +1,120 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import httpx
+from .core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
+from .flowservice.client import AsyncFlowserviceClient, FlowserviceClient
+from .records.client import AsyncRecordsClient, RecordsClient
+
+
+class SkyflowAuth:
+ """
+ Use this class to access the different functions within the SDK. You can instantiate any number of clients with different configuration that will propagate to these functions.
+
+ Parameters
+ ----------
+ base_url : str
+ The base url to use for requests from the client.
+
+ headers : typing.Optional[typing.Dict[str, str]]
+ Additional headers to send with every request.
+
+ timeout : typing.Optional[float]
+ The timeout to be used, in seconds, for requests. By default the timeout is 60 seconds, unless a custom httpx client is used, in which case this default is not enforced.
+
+ follow_redirects : typing.Optional[bool]
+ Whether the default httpx client follows redirects or not, this is irrelevant if a custom httpx client is passed in.
+
+ httpx_client : typing.Optional[httpx.Client]
+ The httpx client to use for making requests, a preconfigured client is used by default, however this is useful should you want to pass in any custom httpx configuration.
+
+ Examples
+ --------
+ from skyflow import SkyflowAuth
+
+ client = SkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+ """
+
+ def __init__(
+ self,
+ *,
+ base_url: str,
+ headers: typing.Optional[typing.Dict[str, str]] = None,
+ timeout: typing.Optional[float] = None,
+ follow_redirects: typing.Optional[bool] = True,
+ httpx_client: typing.Optional[httpx.Client] = None,
+ ):
+ _defaulted_timeout = (
+ timeout if timeout is not None else 60 if httpx_client is None else httpx_client.timeout.read
+ )
+ self._client_wrapper = SyncClientWrapper(
+ base_url=base_url,
+ headers=headers,
+ httpx_client=httpx_client
+ if httpx_client is not None
+ else httpx.Client(timeout=_defaulted_timeout, follow_redirects=follow_redirects)
+ if follow_redirects is not None
+ else httpx.Client(timeout=_defaulted_timeout),
+ timeout=_defaulted_timeout,
+ )
+ self.records = RecordsClient(client_wrapper=self._client_wrapper)
+ self.flowservice = FlowserviceClient(client_wrapper=self._client_wrapper)
+
+
+class AsyncSkyflowAuth:
+ """
+ Use this class to access the different functions within the SDK. You can instantiate any number of clients with different configuration that will propagate to these functions.
+
+ Parameters
+ ----------
+ base_url : str
+ The base url to use for requests from the client.
+
+ headers : typing.Optional[typing.Dict[str, str]]
+ Additional headers to send with every request.
+
+ timeout : typing.Optional[float]
+ The timeout to be used, in seconds, for requests. By default the timeout is 60 seconds, unless a custom httpx client is used, in which case this default is not enforced.
+
+ follow_redirects : typing.Optional[bool]
+ Whether the default httpx client follows redirects or not, this is irrelevant if a custom httpx client is passed in.
+
+ httpx_client : typing.Optional[httpx.AsyncClient]
+ The httpx client to use for making requests, a preconfigured client is used by default, however this is useful should you want to pass in any custom httpx configuration.
+
+ Examples
+ --------
+ from skyflow import AsyncSkyflowAuth
+
+ client = AsyncSkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+ """
+
+ def __init__(
+ self,
+ *,
+ base_url: str,
+ headers: typing.Optional[typing.Dict[str, str]] = None,
+ timeout: typing.Optional[float] = None,
+ follow_redirects: typing.Optional[bool] = True,
+ httpx_client: typing.Optional[httpx.AsyncClient] = None,
+ ):
+ _defaulted_timeout = (
+ timeout if timeout is not None else 60 if httpx_client is None else httpx_client.timeout.read
+ )
+ self._client_wrapper = AsyncClientWrapper(
+ base_url=base_url,
+ headers=headers,
+ httpx_client=httpx_client
+ if httpx_client is not None
+ else httpx.AsyncClient(timeout=_defaulted_timeout, follow_redirects=follow_redirects)
+ if follow_redirects is not None
+ else httpx.AsyncClient(timeout=_defaulted_timeout),
+ timeout=_defaulted_timeout,
+ )
+ self.records = AsyncRecordsClient(client_wrapper=self._client_wrapper)
+ self.flowservice = AsyncFlowserviceClient(client_wrapper=self._client_wrapper)
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/__init__.py b/flowvault/skyflow_flowvault/generated/rest/core/__init__.py
new file mode 100644
index 00000000..31bbb818
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/__init__.py
@@ -0,0 +1,52 @@
+# This file was auto-generated by Fern from our API Definition.
+
+# isort: skip_file
+
+from .api_error import ApiError
+from .client_wrapper import AsyncClientWrapper, BaseClientWrapper, SyncClientWrapper
+from .datetime_utils import serialize_datetime
+from .file import File, convert_file_dict_to_httpx_tuples, with_content_type
+from .http_client import AsyncHttpClient, HttpClient
+from .http_response import AsyncHttpResponse, HttpResponse
+from .jsonable_encoder import jsonable_encoder
+from .pydantic_utilities import (
+ IS_PYDANTIC_V2,
+ UniversalBaseModel,
+ UniversalRootModel,
+ parse_obj_as,
+ universal_field_validator,
+ universal_root_validator,
+ update_forward_refs,
+)
+from .query_encoder import encode_query
+from .remove_none_from_dict import remove_none_from_dict
+from .request_options import RequestOptions
+from .serialization import FieldMetadata, convert_and_respect_annotation_metadata
+
+__all__ = [
+ "ApiError",
+ "AsyncClientWrapper",
+ "AsyncHttpClient",
+ "AsyncHttpResponse",
+ "BaseClientWrapper",
+ "FieldMetadata",
+ "File",
+ "HttpClient",
+ "HttpResponse",
+ "IS_PYDANTIC_V2",
+ "RequestOptions",
+ "SyncClientWrapper",
+ "UniversalBaseModel",
+ "UniversalRootModel",
+ "convert_and_respect_annotation_metadata",
+ "convert_file_dict_to_httpx_tuples",
+ "encode_query",
+ "jsonable_encoder",
+ "parse_obj_as",
+ "remove_none_from_dict",
+ "serialize_datetime",
+ "universal_field_validator",
+ "universal_root_validator",
+ "update_forward_refs",
+ "with_content_type",
+]
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/api_error.py b/flowvault/skyflow_flowvault/generated/rest/core/api_error.py
new file mode 100644
index 00000000..6f850a60
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/api_error.py
@@ -0,0 +1,23 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from typing import Any, Dict, Optional
+
+
+class ApiError(Exception):
+ headers: Optional[Dict[str, str]]
+ status_code: Optional[int]
+ body: Any
+
+ def __init__(
+ self,
+ *,
+ headers: Optional[Dict[str, str]] = None,
+ status_code: Optional[int] = None,
+ body: Any = None,
+ ) -> None:
+ self.headers = headers
+ self.status_code = status_code
+ self.body = body
+
+ def __str__(self) -> str:
+ return f"headers: {self.headers}, status_code: {self.status_code}, body: {self.body}"
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/client_wrapper.py b/flowvault/skyflow_flowvault/generated/rest/core/client_wrapper.py
new file mode 100644
index 00000000..8f63f6ee
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/client_wrapper.py
@@ -0,0 +1,73 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import httpx
+from .http_client import AsyncHttpClient, HttpClient
+
+
+class BaseClientWrapper:
+ def __init__(
+ self,
+ *,
+ headers: typing.Optional[typing.Dict[str, str]] = None,
+ base_url: str,
+ timeout: typing.Optional[float] = None,
+ ):
+ self._headers = headers
+ self._base_url = base_url
+ self._timeout = timeout
+
+ def get_headers(self) -> typing.Dict[str, str]:
+ headers: typing.Dict[str, str] = {
+ "X-Fern-Language": "Python",
+ "X-Fern-SDK-Name": "skyflow.generated.rest",
+ "X-Fern-SDK-Version": "0.0.10",
+ **(self.get_custom_headers() or {}),
+ }
+ return headers
+
+ def get_custom_headers(self) -> typing.Optional[typing.Dict[str, str]]:
+ return self._headers
+
+ def get_base_url(self) -> str:
+ return self._base_url
+
+ def get_timeout(self) -> typing.Optional[float]:
+ return self._timeout
+
+
+class SyncClientWrapper(BaseClientWrapper):
+ def __init__(
+ self,
+ *,
+ headers: typing.Optional[typing.Dict[str, str]] = None,
+ base_url: str,
+ timeout: typing.Optional[float] = None,
+ httpx_client: httpx.Client,
+ ):
+ super().__init__(headers=headers, base_url=base_url, timeout=timeout)
+ self.httpx_client = HttpClient(
+ httpx_client=httpx_client,
+ base_headers=self.get_headers,
+ base_timeout=self.get_timeout,
+ base_url=self.get_base_url,
+ )
+
+
+class AsyncClientWrapper(BaseClientWrapper):
+ def __init__(
+ self,
+ *,
+ headers: typing.Optional[typing.Dict[str, str]] = None,
+ base_url: str,
+ timeout: typing.Optional[float] = None,
+ httpx_client: httpx.AsyncClient,
+ ):
+ super().__init__(headers=headers, base_url=base_url, timeout=timeout)
+ self.httpx_client = AsyncHttpClient(
+ httpx_client=httpx_client,
+ base_headers=self.get_headers,
+ base_timeout=self.get_timeout,
+ base_url=self.get_base_url,
+ )
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/datetime_utils.py b/flowvault/skyflow_flowvault/generated/rest/core/datetime_utils.py
new file mode 100644
index 00000000..7c9864a9
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/datetime_utils.py
@@ -0,0 +1,28 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import datetime as dt
+
+
+def serialize_datetime(v: dt.datetime) -> str:
+ """
+ Serialize a datetime including timezone info.
+
+ Uses the timezone info provided if present, otherwise uses the current runtime's timezone info.
+
+ UTC datetimes end in "Z" while all other timezones are represented as offset from UTC, e.g. +05:00.
+ """
+
+ def _serialize_zoned_datetime(v: dt.datetime) -> str:
+ if v.tzinfo is not None and v.tzinfo.tzname(None) == dt.timezone.utc.tzname(None):
+ # UTC is a special case where we use "Z" at the end instead of "+00:00"
+ return v.isoformat().replace("+00:00", "Z")
+ else:
+ # Delegate to the typical +/- offset format
+ return v.isoformat()
+
+ if v.tzinfo is not None:
+ return _serialize_zoned_datetime(v)
+ else:
+ local_tz = dt.datetime.now().astimezone().tzinfo
+ localized_dt = v.replace(tzinfo=local_tz)
+ return _serialize_zoned_datetime(localized_dt)
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/file.py b/flowvault/skyflow_flowvault/generated/rest/core/file.py
new file mode 100644
index 00000000..44b0d27c
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/file.py
@@ -0,0 +1,67 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from typing import IO, Dict, List, Mapping, Optional, Tuple, Union, cast
+
+# File typing inspired by the flexibility of types within the httpx library
+# https://github.com/encode/httpx/blob/master/httpx/_types.py
+FileContent = Union[IO[bytes], bytes, str]
+File = Union[
+ # file (or bytes)
+ FileContent,
+ # (filename, file (or bytes))
+ Tuple[Optional[str], FileContent],
+ # (filename, file (or bytes), content_type)
+ Tuple[Optional[str], FileContent, Optional[str]],
+ # (filename, file (or bytes), content_type, headers)
+ Tuple[
+ Optional[str],
+ FileContent,
+ Optional[str],
+ Mapping[str, str],
+ ],
+]
+
+
+def convert_file_dict_to_httpx_tuples(
+ d: Dict[str, Union[File, List[File]]],
+) -> List[Tuple[str, File]]:
+ """
+ The format we use is a list of tuples, where the first element is the
+ name of the file and the second is the file object. Typically HTTPX wants
+ a dict, but to be able to send lists of files, you have to use the list
+ approach (which also works for non-lists)
+ https://github.com/encode/httpx/pull/1032
+ """
+
+ httpx_tuples = []
+ for key, file_like in d.items():
+ if isinstance(file_like, list):
+ for file_like_item in file_like:
+ httpx_tuples.append((key, file_like_item))
+ else:
+ httpx_tuples.append((key, file_like))
+ return httpx_tuples
+
+
+def with_content_type(*, file: File, default_content_type: str) -> File:
+ """
+ This function resolves to the file's content type, if provided, and defaults
+ to the default_content_type value if not.
+ """
+ if isinstance(file, tuple):
+ if len(file) == 2:
+ filename, content = cast(Tuple[Optional[str], FileContent], file) # type: ignore
+ return (filename, content, default_content_type)
+ elif len(file) == 3:
+ filename, content, file_content_type = cast(Tuple[Optional[str], FileContent, Optional[str]], file) # type: ignore
+ out_content_type = file_content_type or default_content_type
+ return (filename, content, out_content_type)
+ elif len(file) == 4:
+ filename, content, file_content_type, headers = cast( # type: ignore
+ Tuple[Optional[str], FileContent, Optional[str], Mapping[str, str]], file
+ )
+ out_content_type = file_content_type or default_content_type
+ return (filename, content, out_content_type, headers)
+ else:
+ raise ValueError(f"Unexpected tuple length: {len(file)}")
+ return (None, file, default_content_type)
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/force_multipart.py b/flowvault/skyflow_flowvault/generated/rest/core/force_multipart.py
new file mode 100644
index 00000000..ae24ccff
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/force_multipart.py
@@ -0,0 +1,16 @@
+# This file was auto-generated by Fern from our API Definition.
+
+
+class ForceMultipartDict(dict):
+ """
+ A dictionary subclass that always evaluates to True in boolean contexts.
+
+ This is used to force multipart/form-data encoding in HTTP requests even when
+ the dictionary is empty, which would normally evaluate to False.
+ """
+
+ def __bool__(self):
+ return True
+
+
+FORCE_MULTIPART = ForceMultipartDict()
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/http_client.py b/flowvault/skyflow_flowvault/generated/rest/core/http_client.py
new file mode 100644
index 00000000..e4173f99
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/http_client.py
@@ -0,0 +1,543 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import asyncio
+import email.utils
+import re
+import time
+import typing
+import urllib.parse
+from contextlib import asynccontextmanager, contextmanager
+from random import random
+
+import httpx
+from .file import File, convert_file_dict_to_httpx_tuples
+from .force_multipart import FORCE_MULTIPART
+from .jsonable_encoder import jsonable_encoder
+from .query_encoder import encode_query
+from .remove_none_from_dict import remove_none_from_dict
+from .request_options import RequestOptions
+from httpx._types import RequestFiles
+
+INITIAL_RETRY_DELAY_SECONDS = 0.5
+MAX_RETRY_DELAY_SECONDS = 10
+MAX_RETRY_DELAY_SECONDS_FROM_HEADER = 30
+
+
+def _parse_retry_after(response_headers: httpx.Headers) -> typing.Optional[float]:
+ """
+ This function parses the `Retry-After` header in a HTTP response and returns the number of seconds to wait.
+
+ Inspired by the urllib3 retry implementation.
+ """
+ retry_after_ms = response_headers.get("retry-after-ms")
+ if retry_after_ms is not None:
+ try:
+ return int(retry_after_ms) / 1000 if retry_after_ms > 0 else 0
+ except Exception:
+ pass
+
+ retry_after = response_headers.get("retry-after")
+ if retry_after is None:
+ return None
+
+ # Attempt to parse the header as an int.
+ if re.match(r"^\s*[0-9]+\s*$", retry_after):
+ seconds = float(retry_after)
+ # Fallback to parsing it as a date.
+ else:
+ retry_date_tuple = email.utils.parsedate_tz(retry_after)
+ if retry_date_tuple is None:
+ return None
+ if retry_date_tuple[9] is None: # Python 2
+ # Assume UTC if no timezone was specified
+ # On Python2.7, parsedate_tz returns None for a timezone offset
+ # instead of 0 if no timezone is given, where mktime_tz treats
+ # a None timezone offset as local time.
+ retry_date_tuple = retry_date_tuple[:9] + (0,) + retry_date_tuple[10:]
+
+ retry_date = email.utils.mktime_tz(retry_date_tuple)
+ seconds = retry_date - time.time()
+
+ if seconds < 0:
+ seconds = 0
+
+ return seconds
+
+
+def _retry_timeout(response: httpx.Response, retries: int) -> float:
+ """
+ Determine the amount of time to wait before retrying a request.
+ This function begins by trying to parse a retry-after header from the response, and then proceeds to use exponential backoff
+ with a jitter to determine the number of seconds to wait.
+ """
+
+ # If the API asks us to wait a certain amount of time (and it's a reasonable amount), just do what it says.
+ retry_after = _parse_retry_after(response.headers)
+ if retry_after is not None and retry_after <= MAX_RETRY_DELAY_SECONDS_FROM_HEADER:
+ return retry_after
+
+ # Apply exponential backoff, capped at MAX_RETRY_DELAY_SECONDS.
+ retry_delay = min(INITIAL_RETRY_DELAY_SECONDS * pow(2.0, retries), MAX_RETRY_DELAY_SECONDS)
+
+ # Add a randomness / jitter to the retry delay to avoid overwhelming the server with retries.
+ timeout = retry_delay * (1 - 0.25 * random())
+ return timeout if timeout >= 0 else 0
+
+
+def _should_retry(response: httpx.Response) -> bool:
+ retryable_400s = [429, 408, 409]
+ return response.status_code >= 500 or response.status_code in retryable_400s
+
+
+def remove_omit_from_dict(
+ original: typing.Dict[str, typing.Optional[typing.Any]],
+ omit: typing.Optional[typing.Any],
+) -> typing.Dict[str, typing.Any]:
+ if omit is None:
+ return original
+ new: typing.Dict[str, typing.Any] = {}
+ for key, value in original.items():
+ if value is not omit:
+ new[key] = value
+ return new
+
+
+def maybe_filter_request_body(
+ data: typing.Optional[typing.Any],
+ request_options: typing.Optional[RequestOptions],
+ omit: typing.Optional[typing.Any],
+) -> typing.Optional[typing.Any]:
+ if data is None:
+ return (
+ jsonable_encoder(request_options.get("additional_body_parameters", {})) or {}
+ if request_options is not None
+ else None
+ )
+ elif not isinstance(data, typing.Mapping):
+ data_content = jsonable_encoder(data)
+ else:
+ data_content = {
+ **(jsonable_encoder(remove_omit_from_dict(data, omit))), # type: ignore
+ **(
+ jsonable_encoder(request_options.get("additional_body_parameters", {})) or {}
+ if request_options is not None
+ else {}
+ ),
+ }
+ return data_content
+
+
+# Abstracted out for testing purposes
+def get_request_body(
+ *,
+ json: typing.Optional[typing.Any],
+ data: typing.Optional[typing.Any],
+ request_options: typing.Optional[RequestOptions],
+ omit: typing.Optional[typing.Any],
+) -> typing.Tuple[typing.Optional[typing.Any], typing.Optional[typing.Any]]:
+ json_body = None
+ data_body = None
+ if data is not None:
+ data_body = maybe_filter_request_body(data, request_options, omit)
+ else:
+ # If both data and json are None, we send json data in the event extra properties are specified
+ json_body = maybe_filter_request_body(json, request_options, omit)
+
+ # If you have an empty JSON body, you should just send None
+ return (json_body if json_body != {} else None), data_body if data_body != {} else None
+
+
+class HttpClient:
+ def __init__(
+ self,
+ *,
+ httpx_client: httpx.Client,
+ base_timeout: typing.Callable[[], typing.Optional[float]],
+ base_headers: typing.Callable[[], typing.Dict[str, str]],
+ base_url: typing.Optional[typing.Callable[[], str]] = None,
+ ):
+ self.base_url = base_url
+ self.base_timeout = base_timeout
+ self.base_headers = base_headers
+ self.httpx_client = httpx_client
+
+ def get_base_url(self, maybe_base_url: typing.Optional[str]) -> str:
+ base_url = maybe_base_url
+ if self.base_url is not None and base_url is None:
+ base_url = self.base_url()
+
+ if base_url is None:
+ raise ValueError("A base_url is required to make this request, please provide one and try again.")
+ return base_url
+
+ def request(
+ self,
+ path: typing.Optional[str] = None,
+ *,
+ method: str,
+ base_url: typing.Optional[str] = None,
+ params: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ json: typing.Optional[typing.Any] = None,
+ data: typing.Optional[typing.Any] = None,
+ content: typing.Optional[typing.Union[bytes, typing.Iterator[bytes], typing.AsyncIterator[bytes]]] = None,
+ files: typing.Optional[
+ typing.Union[
+ typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]],
+ typing.List[typing.Tuple[str, File]],
+ ]
+ ] = None,
+ headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ request_options: typing.Optional[RequestOptions] = None,
+ retries: int = 2,
+ omit: typing.Optional[typing.Any] = None,
+ force_multipart: typing.Optional[bool] = None,
+ ) -> httpx.Response:
+ base_url = self.get_base_url(base_url)
+ timeout = (
+ request_options.get("timeout_in_seconds")
+ if request_options is not None and request_options.get("timeout_in_seconds") is not None
+ else self.base_timeout()
+ )
+
+ json_body, data_body = get_request_body(json=json, data=data, request_options=request_options, omit=omit)
+
+ request_files: typing.Optional[RequestFiles] = (
+ convert_file_dict_to_httpx_tuples(remove_omit_from_dict(remove_none_from_dict(files), omit))
+ if (files is not None and files is not omit and isinstance(files, dict))
+ else None
+ )
+
+ if (request_files is None or len(request_files) == 0) and force_multipart:
+ request_files = FORCE_MULTIPART
+
+ response = self.httpx_client.request(
+ method=method,
+ url=urllib.parse.urljoin(f"{base_url}/", path),
+ headers=jsonable_encoder(
+ remove_none_from_dict(
+ {
+ **self.base_headers(),
+ **(headers if headers is not None else {}),
+ **(request_options.get("additional_headers", {}) or {} if request_options is not None else {}),
+ }
+ )
+ ),
+ params=encode_query(
+ jsonable_encoder(
+ remove_none_from_dict(
+ remove_omit_from_dict(
+ {
+ **(params if params is not None else {}),
+ **(
+ request_options.get("additional_query_parameters", {}) or {}
+ if request_options is not None
+ else {}
+ ),
+ },
+ omit,
+ )
+ )
+ )
+ ),
+ json=json_body,
+ data=data_body,
+ content=content,
+ files=request_files,
+ timeout=timeout,
+ )
+
+ max_retries: int = request_options.get("max_retries", 0) if request_options is not None else 0
+ if _should_retry(response=response):
+ if max_retries > retries:
+ time.sleep(_retry_timeout(response=response, retries=retries))
+ return self.request(
+ path=path,
+ method=method,
+ base_url=base_url,
+ params=params,
+ json=json,
+ content=content,
+ files=files,
+ headers=headers,
+ request_options=request_options,
+ retries=retries + 1,
+ omit=omit,
+ )
+
+ return response
+
+ @contextmanager
+ def stream(
+ self,
+ path: typing.Optional[str] = None,
+ *,
+ method: str,
+ base_url: typing.Optional[str] = None,
+ params: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ json: typing.Optional[typing.Any] = None,
+ data: typing.Optional[typing.Any] = None,
+ content: typing.Optional[typing.Union[bytes, typing.Iterator[bytes], typing.AsyncIterator[bytes]]] = None,
+ files: typing.Optional[
+ typing.Union[
+ typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]],
+ typing.List[typing.Tuple[str, File]],
+ ]
+ ] = None,
+ headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ request_options: typing.Optional[RequestOptions] = None,
+ retries: int = 2,
+ omit: typing.Optional[typing.Any] = None,
+ force_multipart: typing.Optional[bool] = None,
+ ) -> typing.Iterator[httpx.Response]:
+ base_url = self.get_base_url(base_url)
+ timeout = (
+ request_options.get("timeout_in_seconds")
+ if request_options is not None and request_options.get("timeout_in_seconds") is not None
+ else self.base_timeout()
+ )
+
+ request_files: typing.Optional[RequestFiles] = (
+ convert_file_dict_to_httpx_tuples(remove_omit_from_dict(remove_none_from_dict(files), omit))
+ if (files is not None and files is not omit and isinstance(files, dict))
+ else None
+ )
+
+ if (request_files is None or len(request_files) == 0) and force_multipart:
+ request_files = FORCE_MULTIPART
+
+ json_body, data_body = get_request_body(json=json, data=data, request_options=request_options, omit=omit)
+
+ with self.httpx_client.stream(
+ method=method,
+ url=urllib.parse.urljoin(f"{base_url}/", path),
+ headers=jsonable_encoder(
+ remove_none_from_dict(
+ {
+ **self.base_headers(),
+ **(headers if headers is not None else {}),
+ **(request_options.get("additional_headers", {}) if request_options is not None else {}),
+ }
+ )
+ ),
+ params=encode_query(
+ jsonable_encoder(
+ remove_none_from_dict(
+ remove_omit_from_dict(
+ {
+ **(params if params is not None else {}),
+ **(
+ request_options.get("additional_query_parameters", {})
+ if request_options is not None
+ else {}
+ ),
+ },
+ omit,
+ )
+ )
+ )
+ ),
+ json=json_body,
+ data=data_body,
+ content=content,
+ files=request_files,
+ timeout=timeout,
+ ) as stream:
+ yield stream
+
+
+class AsyncHttpClient:
+ def __init__(
+ self,
+ *,
+ httpx_client: httpx.AsyncClient,
+ base_timeout: typing.Callable[[], typing.Optional[float]],
+ base_headers: typing.Callable[[], typing.Dict[str, str]],
+ base_url: typing.Optional[typing.Callable[[], str]] = None,
+ ):
+ self.base_url = base_url
+ self.base_timeout = base_timeout
+ self.base_headers = base_headers
+ self.httpx_client = httpx_client
+
+ def get_base_url(self, maybe_base_url: typing.Optional[str]) -> str:
+ base_url = maybe_base_url
+ if self.base_url is not None and base_url is None:
+ base_url = self.base_url()
+
+ if base_url is None:
+ raise ValueError("A base_url is required to make this request, please provide one and try again.")
+ return base_url
+
+ async def request(
+ self,
+ path: typing.Optional[str] = None,
+ *,
+ method: str,
+ base_url: typing.Optional[str] = None,
+ params: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ json: typing.Optional[typing.Any] = None,
+ data: typing.Optional[typing.Any] = None,
+ content: typing.Optional[typing.Union[bytes, typing.Iterator[bytes], typing.AsyncIterator[bytes]]] = None,
+ files: typing.Optional[
+ typing.Union[
+ typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]],
+ typing.List[typing.Tuple[str, File]],
+ ]
+ ] = None,
+ headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ request_options: typing.Optional[RequestOptions] = None,
+ retries: int = 2,
+ omit: typing.Optional[typing.Any] = None,
+ force_multipart: typing.Optional[bool] = None,
+ ) -> httpx.Response:
+ base_url = self.get_base_url(base_url)
+ timeout = (
+ request_options.get("timeout_in_seconds")
+ if request_options is not None and request_options.get("timeout_in_seconds") is not None
+ else self.base_timeout()
+ )
+
+ request_files: typing.Optional[RequestFiles] = (
+ convert_file_dict_to_httpx_tuples(remove_omit_from_dict(remove_none_from_dict(files), omit))
+ if (files is not None and files is not omit and isinstance(files, dict))
+ else None
+ )
+
+ if (request_files is None or len(request_files) == 0) and force_multipart:
+ request_files = FORCE_MULTIPART
+
+ json_body, data_body = get_request_body(json=json, data=data, request_options=request_options, omit=omit)
+
+ # Add the input to each of these and do None-safety checks
+ response = await self.httpx_client.request(
+ method=method,
+ url=urllib.parse.urljoin(f"{base_url}/", path),
+ headers=jsonable_encoder(
+ remove_none_from_dict(
+ {
+ **self.base_headers(),
+ **(headers if headers is not None else {}),
+ **(request_options.get("additional_headers", {}) or {} if request_options is not None else {}),
+ }
+ )
+ ),
+ params=encode_query(
+ jsonable_encoder(
+ remove_none_from_dict(
+ remove_omit_from_dict(
+ {
+ **(params if params is not None else {}),
+ **(
+ request_options.get("additional_query_parameters", {}) or {}
+ if request_options is not None
+ else {}
+ ),
+ },
+ omit,
+ )
+ )
+ )
+ ),
+ json=json_body,
+ data=data_body,
+ content=content,
+ files=request_files,
+ timeout=timeout,
+ )
+
+ max_retries: int = request_options.get("max_retries", 0) if request_options is not None else 0
+ if _should_retry(response=response):
+ if max_retries > retries:
+ await asyncio.sleep(_retry_timeout(response=response, retries=retries))
+ return await self.request(
+ path=path,
+ method=method,
+ base_url=base_url,
+ params=params,
+ json=json,
+ content=content,
+ files=files,
+ headers=headers,
+ request_options=request_options,
+ retries=retries + 1,
+ omit=omit,
+ )
+ return response
+
+ @asynccontextmanager
+ async def stream(
+ self,
+ path: typing.Optional[str] = None,
+ *,
+ method: str,
+ base_url: typing.Optional[str] = None,
+ params: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ json: typing.Optional[typing.Any] = None,
+ data: typing.Optional[typing.Any] = None,
+ content: typing.Optional[typing.Union[bytes, typing.Iterator[bytes], typing.AsyncIterator[bytes]]] = None,
+ files: typing.Optional[
+ typing.Union[
+ typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]],
+ typing.List[typing.Tuple[str, File]],
+ ]
+ ] = None,
+ headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ request_options: typing.Optional[RequestOptions] = None,
+ retries: int = 2,
+ omit: typing.Optional[typing.Any] = None,
+ force_multipart: typing.Optional[bool] = None,
+ ) -> typing.AsyncIterator[httpx.Response]:
+ base_url = self.get_base_url(base_url)
+ timeout = (
+ request_options.get("timeout_in_seconds")
+ if request_options is not None and request_options.get("timeout_in_seconds") is not None
+ else self.base_timeout()
+ )
+
+ request_files: typing.Optional[RequestFiles] = (
+ convert_file_dict_to_httpx_tuples(remove_omit_from_dict(remove_none_from_dict(files), omit))
+ if (files is not None and files is not omit and isinstance(files, dict))
+ else None
+ )
+
+ if (request_files is None or len(request_files) == 0) and force_multipart:
+ request_files = FORCE_MULTIPART
+
+ json_body, data_body = get_request_body(json=json, data=data, request_options=request_options, omit=omit)
+
+ async with self.httpx_client.stream(
+ method=method,
+ url=urllib.parse.urljoin(f"{base_url}/", path),
+ headers=jsonable_encoder(
+ remove_none_from_dict(
+ {
+ **self.base_headers(),
+ **(headers if headers is not None else {}),
+ **(request_options.get("additional_headers", {}) if request_options is not None else {}),
+ }
+ )
+ ),
+ params=encode_query(
+ jsonable_encoder(
+ remove_none_from_dict(
+ remove_omit_from_dict(
+ {
+ **(params if params is not None else {}),
+ **(
+ request_options.get("additional_query_parameters", {})
+ if request_options is not None
+ else {}
+ ),
+ },
+ omit=omit,
+ )
+ )
+ )
+ ),
+ json=json_body,
+ data=data_body,
+ content=content,
+ files=request_files,
+ timeout=timeout,
+ ) as stream:
+ yield stream
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/http_response.py b/flowvault/skyflow_flowvault/generated/rest/core/http_response.py
new file mode 100644
index 00000000..48a1798a
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/http_response.py
@@ -0,0 +1,55 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from typing import Dict, Generic, TypeVar
+
+import httpx
+
+T = TypeVar("T")
+"""Generic to represent the underlying type of the data wrapped by the HTTP response."""
+
+
+class BaseHttpResponse:
+ """Minimalist HTTP response wrapper that exposes response headers."""
+
+ _response: httpx.Response
+
+ def __init__(self, response: httpx.Response):
+ self._response = response
+
+ @property
+ def headers(self) -> Dict[str, str]:
+ return dict(self._response.headers)
+
+
+class HttpResponse(Generic[T], BaseHttpResponse):
+ """HTTP response wrapper that exposes response headers and data."""
+
+ _data: T
+
+ def __init__(self, response: httpx.Response, data: T):
+ super().__init__(response)
+ self._data = data
+
+ @property
+ def data(self) -> T:
+ return self._data
+
+ def close(self) -> None:
+ self._response.close()
+
+
+class AsyncHttpResponse(Generic[T], BaseHttpResponse):
+ """HTTP response wrapper that exposes response headers and data."""
+
+ _data: T
+
+ def __init__(self, response: httpx.Response, data: T):
+ super().__init__(response)
+ self._data = data
+
+ @property
+ def data(self) -> T:
+ return self._data
+
+ async def close(self) -> None:
+ await self._response.aclose()
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/jsonable_encoder.py b/flowvault/skyflow_flowvault/generated/rest/core/jsonable_encoder.py
new file mode 100644
index 00000000..afee3662
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/jsonable_encoder.py
@@ -0,0 +1,100 @@
+# This file was auto-generated by Fern from our API Definition.
+
+"""
+jsonable_encoder converts a Python object to a JSON-friendly dict
+(e.g. datetimes to strings, Pydantic models to dicts).
+
+Taken from FastAPI, and made a bit simpler
+https://github.com/tiangolo/fastapi/blob/master/fastapi/encoders.py
+"""
+
+import base64
+import dataclasses
+import datetime as dt
+from enum import Enum
+from pathlib import PurePath
+from types import GeneratorType
+from typing import Any, Callable, Dict, List, Optional, Set, Union
+
+import pydantic
+from .datetime_utils import serialize_datetime
+from .pydantic_utilities import (
+ IS_PYDANTIC_V2,
+ encode_by_type,
+ to_jsonable_with_fallback,
+)
+
+SetIntStr = Set[Union[int, str]]
+DictIntStrAny = Dict[Union[int, str], Any]
+
+
+def jsonable_encoder(obj: Any, custom_encoder: Optional[Dict[Any, Callable[[Any], Any]]] = None) -> Any:
+ custom_encoder = custom_encoder or {}
+ if custom_encoder:
+ if type(obj) in custom_encoder:
+ return custom_encoder[type(obj)](obj)
+ else:
+ for encoder_type, encoder_instance in custom_encoder.items():
+ if isinstance(obj, encoder_type):
+ return encoder_instance(obj)
+ if isinstance(obj, pydantic.BaseModel):
+ if IS_PYDANTIC_V2:
+ encoder = getattr(obj.model_config, "json_encoders", {}) # type: ignore # Pydantic v2
+ else:
+ encoder = getattr(obj.__config__, "json_encoders", {}) # type: ignore # Pydantic v1
+ if custom_encoder:
+ encoder.update(custom_encoder)
+ obj_dict = obj.dict(by_alias=True)
+ if "__root__" in obj_dict:
+ obj_dict = obj_dict["__root__"]
+ if "root" in obj_dict:
+ obj_dict = obj_dict["root"]
+ return jsonable_encoder(obj_dict, custom_encoder=encoder)
+ if dataclasses.is_dataclass(obj):
+ obj_dict = dataclasses.asdict(obj) # type: ignore
+ return jsonable_encoder(obj_dict, custom_encoder=custom_encoder)
+ if isinstance(obj, bytes):
+ return base64.b64encode(obj).decode("utf-8")
+ if isinstance(obj, Enum):
+ return obj.value
+ if isinstance(obj, PurePath):
+ return str(obj)
+ if isinstance(obj, (str, int, float, type(None))):
+ return obj
+ if isinstance(obj, dt.datetime):
+ return serialize_datetime(obj)
+ if isinstance(obj, dt.date):
+ return str(obj)
+ if isinstance(obj, dict):
+ encoded_dict = {}
+ allowed_keys = set(obj.keys())
+ for key, value in obj.items():
+ if key in allowed_keys:
+ encoded_key = jsonable_encoder(key, custom_encoder=custom_encoder)
+ encoded_value = jsonable_encoder(value, custom_encoder=custom_encoder)
+ encoded_dict[encoded_key] = encoded_value
+ return encoded_dict
+ if isinstance(obj, (list, set, frozenset, GeneratorType, tuple)):
+ encoded_list = []
+ for item in obj:
+ encoded_list.append(jsonable_encoder(item, custom_encoder=custom_encoder))
+ return encoded_list
+
+ def fallback_serializer(o: Any) -> Any:
+ attempt_encode = encode_by_type(o)
+ if attempt_encode is not None:
+ return attempt_encode
+
+ try:
+ data = dict(o)
+ except Exception as e:
+ errors: List[Exception] = []
+ errors.append(e)
+ try:
+ data = vars(o)
+ except Exception as e:
+ errors.append(e)
+ raise ValueError(errors) from e
+ return jsonable_encoder(data, custom_encoder=custom_encoder)
+
+ return to_jsonable_with_fallback(obj, fallback_serializer)
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/pydantic_utilities.py b/flowvault/skyflow_flowvault/generated/rest/core/pydantic_utilities.py
new file mode 100644
index 00000000..7db29500
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/pydantic_utilities.py
@@ -0,0 +1,255 @@
+# This file was auto-generated by Fern from our API Definition.
+
+# nopycln: file
+import datetime as dt
+from collections import defaultdict
+from typing import Any, Callable, ClassVar, Dict, List, Mapping, Optional, Set, Tuple, Type, TypeVar, Union, cast
+
+import pydantic
+
+IS_PYDANTIC_V2 = pydantic.VERSION.startswith("2.")
+
+if IS_PYDANTIC_V2:
+ from pydantic.v1.datetime_parse import parse_date as parse_date
+ from pydantic.v1.datetime_parse import parse_datetime as parse_datetime
+ from pydantic.v1.fields import ModelField as ModelField
+ from pydantic.v1.json import ENCODERS_BY_TYPE as encoders_by_type # type: ignore[attr-defined]
+ from pydantic.v1.typing import get_args as get_args
+ from pydantic.v1.typing import get_origin as get_origin
+ from pydantic.v1.typing import is_literal_type as is_literal_type
+ from pydantic.v1.typing import is_union as is_union
+else:
+ from pydantic.datetime_parse import parse_date as parse_date # type: ignore[no-redef]
+ from pydantic.datetime_parse import parse_datetime as parse_datetime # type: ignore[no-redef]
+ from pydantic.fields import ModelField as ModelField # type: ignore[attr-defined, no-redef]
+ from pydantic.json import ENCODERS_BY_TYPE as encoders_by_type # type: ignore[no-redef]
+ from pydantic.typing import get_args as get_args # type: ignore[no-redef]
+ from pydantic.typing import get_origin as get_origin # type: ignore[no-redef]
+ from pydantic.typing import is_literal_type as is_literal_type # type: ignore[no-redef]
+ from pydantic.typing import is_union as is_union # type: ignore[no-redef]
+
+from .datetime_utils import serialize_datetime
+from .serialization import convert_and_respect_annotation_metadata
+from typing_extensions import TypeAlias
+
+T = TypeVar("T")
+Model = TypeVar("Model", bound=pydantic.BaseModel)
+
+
+def parse_obj_as(type_: Type[T], object_: Any) -> T:
+ dealiased_object = convert_and_respect_annotation_metadata(object_=object_, annotation=type_, direction="read")
+ if IS_PYDANTIC_V2:
+ adapter = pydantic.TypeAdapter(type_) # type: ignore[attr-defined]
+ return adapter.validate_python(dealiased_object)
+ return pydantic.parse_obj_as(type_, dealiased_object)
+
+
+def to_jsonable_with_fallback(obj: Any, fallback_serializer: Callable[[Any], Any]) -> Any:
+ if IS_PYDANTIC_V2:
+ from pydantic_core import to_jsonable_python
+
+ return to_jsonable_python(obj, fallback=fallback_serializer)
+ return fallback_serializer(obj)
+
+
+class UniversalBaseModel(pydantic.BaseModel):
+ if IS_PYDANTIC_V2:
+ model_config: ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict( # type: ignore[typeddict-unknown-key]
+ # Allow fields beginning with `model_` to be used in the model
+ protected_namespaces=(),
+ )
+
+ @pydantic.model_serializer(mode="plain", when_used="json") # type: ignore[attr-defined]
+ def serialize_model(self) -> Any: # type: ignore[name-defined]
+ serialized = self.model_dump()
+ data = {k: serialize_datetime(v) if isinstance(v, dt.datetime) else v for k, v in serialized.items()}
+ return data
+
+ else:
+
+ class Config:
+ smart_union = True
+ json_encoders = {dt.datetime: serialize_datetime}
+
+ @classmethod
+ def model_construct(cls: Type["Model"], _fields_set: Optional[Set[str]] = None, **values: Any) -> "Model":
+ dealiased_object = convert_and_respect_annotation_metadata(object_=values, annotation=cls, direction="read")
+ return cls.construct(_fields_set, **dealiased_object)
+
+ @classmethod
+ def construct(cls: Type["Model"], _fields_set: Optional[Set[str]] = None, **values: Any) -> "Model":
+ dealiased_object = convert_and_respect_annotation_metadata(object_=values, annotation=cls, direction="read")
+ if IS_PYDANTIC_V2:
+ return super().model_construct(_fields_set, **dealiased_object) # type: ignore[misc]
+ return super().construct(_fields_set, **dealiased_object)
+
+ def json(self, **kwargs: Any) -> str:
+ kwargs_with_defaults = {
+ "by_alias": True,
+ "exclude_unset": True,
+ **kwargs,
+ }
+ if IS_PYDANTIC_V2:
+ return super().model_dump_json(**kwargs_with_defaults) # type: ignore[misc]
+ return super().json(**kwargs_with_defaults)
+
+ def dict(self, **kwargs: Any) -> Dict[str, Any]:
+ """
+ Override the default dict method to `exclude_unset` by default. This function patches
+ `exclude_unset` to work include fields within non-None default values.
+ """
+ # Note: the logic here is multiplexed given the levers exposed in Pydantic V1 vs V2
+ # Pydantic V1's .dict can be extremely slow, so we do not want to call it twice.
+ #
+ # We'd ideally do the same for Pydantic V2, but it shells out to a library to serialize models
+ # that we have less control over, and this is less intrusive than custom serializers for now.
+ if IS_PYDANTIC_V2:
+ kwargs_with_defaults_exclude_unset = {
+ **kwargs,
+ "by_alias": True,
+ "exclude_unset": True,
+ "exclude_none": False,
+ }
+ kwargs_with_defaults_exclude_none = {
+ **kwargs,
+ "by_alias": True,
+ "exclude_none": True,
+ "exclude_unset": False,
+ }
+ dict_dump = deep_union_pydantic_dicts(
+ super().model_dump(**kwargs_with_defaults_exclude_unset), # type: ignore[misc]
+ super().model_dump(**kwargs_with_defaults_exclude_none), # type: ignore[misc]
+ )
+
+ else:
+ _fields_set = self.__fields_set__.copy()
+
+ fields = _get_model_fields(self.__class__)
+ for name, field in fields.items():
+ if name not in _fields_set:
+ default = _get_field_default(field)
+
+ # If the default values are non-null act like they've been set
+ # This effectively allows exclude_unset to work like exclude_none where
+ # the latter passes through intentionally set none values.
+ if default is not None or ("exclude_unset" in kwargs and not kwargs["exclude_unset"]):
+ _fields_set.add(name)
+
+ if default is not None:
+ self.__fields_set__.add(name)
+
+ kwargs_with_defaults_exclude_unset_include_fields = {
+ "by_alias": True,
+ "exclude_unset": True,
+ "include": _fields_set,
+ **kwargs,
+ }
+
+ dict_dump = super().dict(**kwargs_with_defaults_exclude_unset_include_fields)
+
+ return convert_and_respect_annotation_metadata(object_=dict_dump, annotation=self.__class__, direction="write")
+
+
+def _union_list_of_pydantic_dicts(source: List[Any], destination: List[Any]) -> List[Any]:
+ converted_list: List[Any] = []
+ for i, item in enumerate(source):
+ destination_value = destination[i]
+ if isinstance(item, dict):
+ converted_list.append(deep_union_pydantic_dicts(item, destination_value))
+ elif isinstance(item, list):
+ converted_list.append(_union_list_of_pydantic_dicts(item, destination_value))
+ else:
+ converted_list.append(item)
+ return converted_list
+
+
+def deep_union_pydantic_dicts(source: Dict[str, Any], destination: Dict[str, Any]) -> Dict[str, Any]:
+ for key, value in source.items():
+ node = destination.setdefault(key, {})
+ if isinstance(value, dict):
+ deep_union_pydantic_dicts(value, node)
+ # Note: we do not do this same processing for sets given we do not have sets of models
+ # and given the sets are unordered, the processing of the set and matching objects would
+ # be non-trivial.
+ elif isinstance(value, list):
+ destination[key] = _union_list_of_pydantic_dicts(value, node)
+ else:
+ destination[key] = value
+
+ return destination
+
+
+if IS_PYDANTIC_V2:
+
+ class V2RootModel(UniversalBaseModel, pydantic.RootModel): # type: ignore[misc, name-defined, type-arg]
+ pass
+
+ UniversalRootModel: TypeAlias = V2RootModel # type: ignore[misc]
+else:
+ UniversalRootModel: TypeAlias = UniversalBaseModel # type: ignore[misc, no-redef]
+
+
+def encode_by_type(o: Any) -> Any:
+ encoders_by_class_tuples: Dict[Callable[[Any], Any], Tuple[Any, ...]] = defaultdict(tuple)
+ for type_, encoder in encoders_by_type.items():
+ encoders_by_class_tuples[encoder] += (type_,)
+
+ if type(o) in encoders_by_type:
+ return encoders_by_type[type(o)](o)
+ for encoder, classes_tuple in encoders_by_class_tuples.items():
+ if isinstance(o, classes_tuple):
+ return encoder(o)
+
+
+def update_forward_refs(model: Type["Model"], **localns: Any) -> None:
+ if IS_PYDANTIC_V2:
+ model.model_rebuild(raise_errors=False) # type: ignore[attr-defined]
+ else:
+ model.update_forward_refs(**localns)
+
+
+# Mirrors Pydantic's internal typing
+AnyCallable = Callable[..., Any]
+
+
+def universal_root_validator(
+ pre: bool = False,
+) -> Callable[[AnyCallable], AnyCallable]:
+ def decorator(func: AnyCallable) -> AnyCallable:
+ if IS_PYDANTIC_V2:
+ return cast(AnyCallable, pydantic.model_validator(mode="before" if pre else "after")(func)) # type: ignore[attr-defined]
+ return cast(AnyCallable, pydantic.root_validator(pre=pre)(func)) # type: ignore[call-overload]
+
+ return decorator
+
+
+def universal_field_validator(field_name: str, pre: bool = False) -> Callable[[AnyCallable], AnyCallable]:
+ def decorator(func: AnyCallable) -> AnyCallable:
+ if IS_PYDANTIC_V2:
+ return cast(AnyCallable, pydantic.field_validator(field_name, mode="before" if pre else "after")(func)) # type: ignore[attr-defined]
+ return cast(AnyCallable, pydantic.validator(field_name, pre=pre)(func))
+
+ return decorator
+
+
+PydanticField = Union[ModelField, pydantic.fields.FieldInfo]
+
+
+def _get_model_fields(model: Type["Model"]) -> Mapping[str, PydanticField]:
+ if IS_PYDANTIC_V2:
+ return cast(Mapping[str, PydanticField], model.model_fields) # type: ignore[attr-defined]
+ return cast(Mapping[str, PydanticField], model.__fields__)
+
+
+def _get_field_default(field: PydanticField) -> Any:
+ try:
+ value = field.get_default() # type: ignore[union-attr]
+ except:
+ value = field.default
+ if IS_PYDANTIC_V2:
+ from pydantic_core import PydanticUndefined
+
+ if value == PydanticUndefined:
+ return None
+ return value
+ return value
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/query_encoder.py b/flowvault/skyflow_flowvault/generated/rest/core/query_encoder.py
new file mode 100644
index 00000000..3183001d
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/query_encoder.py
@@ -0,0 +1,58 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import pydantic
+
+
+# Flattens dicts to be of the form {"key[subkey][subkey2]": value} where value is not a dict
+def traverse_query_dict(dict_flat: Dict[str, Any], key_prefix: Optional[str] = None) -> List[Tuple[str, Any]]:
+ result = []
+ for k, v in dict_flat.items():
+ key = f"{key_prefix}[{k}]" if key_prefix is not None else k
+ if isinstance(v, dict):
+ result.extend(traverse_query_dict(v, key))
+ elif isinstance(v, list):
+ for arr_v in v:
+ if isinstance(arr_v, dict):
+ result.extend(traverse_query_dict(arr_v, key))
+ else:
+ result.append((key, arr_v))
+ else:
+ result.append((key, v))
+ return result
+
+
+def single_query_encoder(query_key: str, query_value: Any) -> List[Tuple[str, Any]]:
+ if isinstance(query_value, pydantic.BaseModel) or isinstance(query_value, dict):
+ if isinstance(query_value, pydantic.BaseModel):
+ obj_dict = query_value.dict(by_alias=True)
+ else:
+ obj_dict = query_value
+ return traverse_query_dict(obj_dict, query_key)
+ elif isinstance(query_value, list):
+ encoded_values: List[Tuple[str, Any]] = []
+ for value in query_value:
+ if isinstance(value, pydantic.BaseModel) or isinstance(value, dict):
+ if isinstance(value, pydantic.BaseModel):
+ obj_dict = value.dict(by_alias=True)
+ elif isinstance(value, dict):
+ obj_dict = value
+
+ encoded_values.extend(single_query_encoder(query_key, obj_dict))
+ else:
+ encoded_values.append((query_key, value))
+
+ return encoded_values
+
+ return [(query_key, query_value)]
+
+
+def encode_query(query: Optional[Dict[str, Any]]) -> Optional[List[Tuple[str, Any]]]:
+ if query is None:
+ return None
+
+ encoded_query = []
+ for k, v in query.items():
+ encoded_query.extend(single_query_encoder(k, v))
+ return encoded_query
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/remove_none_from_dict.py b/flowvault/skyflow_flowvault/generated/rest/core/remove_none_from_dict.py
new file mode 100644
index 00000000..c2298143
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/remove_none_from_dict.py
@@ -0,0 +1,11 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from typing import Any, Dict, Mapping, Optional
+
+
+def remove_none_from_dict(original: Mapping[str, Optional[Any]]) -> Dict[str, Any]:
+ new: Dict[str, Any] = {}
+ for key, value in original.items():
+ if value is not None:
+ new[key] = value
+ return new
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/request_options.py b/flowvault/skyflow_flowvault/generated/rest/core/request_options.py
new file mode 100644
index 00000000..1b388044
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/request_options.py
@@ -0,0 +1,35 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+try:
+ from typing import NotRequired # type: ignore
+except ImportError:
+ from typing_extensions import NotRequired
+
+
+class RequestOptions(typing.TypedDict, total=False):
+ """
+ Additional options for request-specific configuration when calling APIs via the SDK.
+ This is used primarily as an optional final parameter for service functions.
+
+ Attributes:
+ - timeout_in_seconds: int. The number of seconds to await an API call before timing out.
+
+ - max_retries: int. The max number of retries to attempt if the API call fails.
+
+ - additional_headers: typing.Dict[str, typing.Any]. A dictionary containing additional parameters to spread into the request's header dict
+
+ - additional_query_parameters: typing.Dict[str, typing.Any]. A dictionary containing additional parameters to spread into the request's query parameters dict
+
+ - additional_body_parameters: typing.Dict[str, typing.Any]. A dictionary containing additional parameters to spread into the request's body parameters dict
+
+ - chunk_size: int. The size, in bytes, to process each chunk of data being streamed back within the response. This equates to leveraging `chunk_size` within `requests` or `httpx`, and is only leveraged for file downloads.
+ """
+
+ timeout_in_seconds: NotRequired[int]
+ max_retries: NotRequired[int]
+ additional_headers: NotRequired[typing.Dict[str, typing.Any]]
+ additional_query_parameters: NotRequired[typing.Dict[str, typing.Any]]
+ additional_body_parameters: NotRequired[typing.Dict[str, typing.Any]]
+ chunk_size: NotRequired[int]
diff --git a/flowvault/skyflow_flowvault/generated/rest/core/serialization.py b/flowvault/skyflow_flowvault/generated/rest/core/serialization.py
new file mode 100644
index 00000000..c36e865c
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/core/serialization.py
@@ -0,0 +1,276 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import collections
+import inspect
+import typing
+
+import pydantic
+import typing_extensions
+
+
+class FieldMetadata:
+ """
+ Metadata class used to annotate fields to provide additional information.
+
+ Example:
+ class MyDict(TypedDict):
+ field: typing.Annotated[str, FieldMetadata(alias="field_name")]
+
+ Will serialize: `{"field": "value"}`
+ To: `{"field_name": "value"}`
+ """
+
+ alias: str
+
+ def __init__(self, *, alias: str) -> None:
+ self.alias = alias
+
+
+def convert_and_respect_annotation_metadata(
+ *,
+ object_: typing.Any,
+ annotation: typing.Any,
+ inner_type: typing.Optional[typing.Any] = None,
+ direction: typing.Literal["read", "write"],
+) -> typing.Any:
+ """
+ Respect the metadata annotations on a field, such as aliasing. This function effectively
+ manipulates the dict-form of an object to respect the metadata annotations. This is primarily used for
+ TypedDicts, which cannot support aliasing out of the box, and can be extended for additional
+ utilities, such as defaults.
+
+ Parameters
+ ----------
+ object_ : typing.Any
+
+ annotation : type
+ The type we're looking to apply typing annotations from
+
+ inner_type : typing.Optional[type]
+
+ Returns
+ -------
+ typing.Any
+ """
+
+ if object_ is None:
+ return None
+ if inner_type is None:
+ inner_type = annotation
+
+ clean_type = _remove_annotations(inner_type)
+ # Pydantic models
+ if (
+ inspect.isclass(clean_type)
+ and issubclass(clean_type, pydantic.BaseModel)
+ and isinstance(object_, typing.Mapping)
+ ):
+ return _convert_mapping(object_, clean_type, direction)
+ # TypedDicts
+ if typing_extensions.is_typeddict(clean_type) and isinstance(object_, typing.Mapping):
+ return _convert_mapping(object_, clean_type, direction)
+
+ if (
+ typing_extensions.get_origin(clean_type) == typing.Dict
+ or typing_extensions.get_origin(clean_type) == dict
+ or clean_type == typing.Dict
+ ) and isinstance(object_, typing.Dict):
+ key_type = typing_extensions.get_args(clean_type)[0]
+ value_type = typing_extensions.get_args(clean_type)[1]
+
+ return {
+ key: convert_and_respect_annotation_metadata(
+ object_=value,
+ annotation=annotation,
+ inner_type=value_type,
+ direction=direction,
+ )
+ for key, value in object_.items()
+ }
+
+ # If you're iterating on a string, do not bother to coerce it to a sequence.
+ if not isinstance(object_, str):
+ if (
+ typing_extensions.get_origin(clean_type) == typing.Set
+ or typing_extensions.get_origin(clean_type) == set
+ or clean_type == typing.Set
+ ) and isinstance(object_, typing.Set):
+ inner_type = typing_extensions.get_args(clean_type)[0]
+ return {
+ convert_and_respect_annotation_metadata(
+ object_=item,
+ annotation=annotation,
+ inner_type=inner_type,
+ direction=direction,
+ )
+ for item in object_
+ }
+ elif (
+ (
+ typing_extensions.get_origin(clean_type) == typing.List
+ or typing_extensions.get_origin(clean_type) == list
+ or clean_type == typing.List
+ )
+ and isinstance(object_, typing.List)
+ ) or (
+ (
+ typing_extensions.get_origin(clean_type) == typing.Sequence
+ or typing_extensions.get_origin(clean_type) == collections.abc.Sequence
+ or clean_type == typing.Sequence
+ )
+ and isinstance(object_, typing.Sequence)
+ ):
+ inner_type = typing_extensions.get_args(clean_type)[0]
+ return [
+ convert_and_respect_annotation_metadata(
+ object_=item,
+ annotation=annotation,
+ inner_type=inner_type,
+ direction=direction,
+ )
+ for item in object_
+ ]
+
+ if typing_extensions.get_origin(clean_type) == typing.Union:
+ # We should be able to ~relatively~ safely try to convert keys against all
+ # member types in the union, the edge case here is if one member aliases a field
+ # of the same name to a different name from another member
+ # Or if another member aliases a field of the same name that another member does not.
+ for member in typing_extensions.get_args(clean_type):
+ object_ = convert_and_respect_annotation_metadata(
+ object_=object_,
+ annotation=annotation,
+ inner_type=member,
+ direction=direction,
+ )
+ return object_
+
+ annotated_type = _get_annotation(annotation)
+ if annotated_type is None:
+ return object_
+
+ # If the object is not a TypedDict, a Union, or other container (list, set, sequence, etc.)
+ # Then we can safely call it on the recursive conversion.
+ return object_
+
+
+def _convert_mapping(
+ object_: typing.Mapping[str, object],
+ expected_type: typing.Any,
+ direction: typing.Literal["read", "write"],
+) -> typing.Mapping[str, object]:
+ converted_object: typing.Dict[str, object] = {}
+ try:
+ annotations = typing_extensions.get_type_hints(expected_type, include_extras=True)
+ except NameError:
+ # The TypedDict contains a circular reference, so
+ # we use the __annotations__ attribute directly.
+ annotations = getattr(expected_type, "__annotations__", {})
+ aliases_to_field_names = _get_alias_to_field_name(annotations)
+ for key, value in object_.items():
+ if direction == "read" and key in aliases_to_field_names:
+ dealiased_key = aliases_to_field_names.get(key)
+ if dealiased_key is not None:
+ type_ = annotations.get(dealiased_key)
+ else:
+ type_ = annotations.get(key)
+ # Note you can't get the annotation by the field name if you're in read mode, so you must check the aliases map
+ #
+ # So this is effectively saying if we're in write mode, and we don't have a type, or if we're in read mode and we don't have an alias
+ # then we can just pass the value through as is
+ if type_ is None:
+ converted_object[key] = value
+ elif direction == "read" and key not in aliases_to_field_names:
+ converted_object[key] = convert_and_respect_annotation_metadata(
+ object_=value, annotation=type_, direction=direction
+ )
+ else:
+ converted_object[_alias_key(key, type_, direction, aliases_to_field_names)] = (
+ convert_and_respect_annotation_metadata(object_=value, annotation=type_, direction=direction)
+ )
+ return converted_object
+
+
+def _get_annotation(type_: typing.Any) -> typing.Optional[typing.Any]:
+ maybe_annotated_type = typing_extensions.get_origin(type_)
+ if maybe_annotated_type is None:
+ return None
+
+ if maybe_annotated_type == typing_extensions.NotRequired:
+ type_ = typing_extensions.get_args(type_)[0]
+ maybe_annotated_type = typing_extensions.get_origin(type_)
+
+ if maybe_annotated_type == typing_extensions.Annotated:
+ return type_
+
+ return None
+
+
+def _remove_annotations(type_: typing.Any) -> typing.Any:
+ maybe_annotated_type = typing_extensions.get_origin(type_)
+ if maybe_annotated_type is None:
+ return type_
+
+ if maybe_annotated_type == typing_extensions.NotRequired:
+ return _remove_annotations(typing_extensions.get_args(type_)[0])
+
+ if maybe_annotated_type == typing_extensions.Annotated:
+ return _remove_annotations(typing_extensions.get_args(type_)[0])
+
+ return type_
+
+
+def get_alias_to_field_mapping(type_: typing.Any) -> typing.Dict[str, str]:
+ annotations = typing_extensions.get_type_hints(type_, include_extras=True)
+ return _get_alias_to_field_name(annotations)
+
+
+def get_field_to_alias_mapping(type_: typing.Any) -> typing.Dict[str, str]:
+ annotations = typing_extensions.get_type_hints(type_, include_extras=True)
+ return _get_field_to_alias_name(annotations)
+
+
+def _get_alias_to_field_name(
+ field_to_hint: typing.Dict[str, typing.Any],
+) -> typing.Dict[str, str]:
+ aliases = {}
+ for field, hint in field_to_hint.items():
+ maybe_alias = _get_alias_from_type(hint)
+ if maybe_alias is not None:
+ aliases[maybe_alias] = field
+ return aliases
+
+
+def _get_field_to_alias_name(
+ field_to_hint: typing.Dict[str, typing.Any],
+) -> typing.Dict[str, str]:
+ aliases = {}
+ for field, hint in field_to_hint.items():
+ maybe_alias = _get_alias_from_type(hint)
+ if maybe_alias is not None:
+ aliases[field] = maybe_alias
+ return aliases
+
+
+def _get_alias_from_type(type_: typing.Any) -> typing.Optional[str]:
+ maybe_annotated_type = _get_annotation(type_)
+
+ if maybe_annotated_type is not None:
+ # The actual annotations are 1 onward, the first is the annotated type
+ annotations = typing_extensions.get_args(maybe_annotated_type)[1:]
+
+ for annotation in annotations:
+ if isinstance(annotation, FieldMetadata) and annotation.alias is not None:
+ return annotation.alias
+ return None
+
+
+def _alias_key(
+ key: str,
+ type_: typing.Any,
+ direction: typing.Literal["read", "write"],
+ aliases_to_field_names: typing.Dict[str, str],
+) -> str:
+ if direction == "read":
+ return aliases_to_field_names.get(key, key)
+ return _get_alias_from_type(type_=type_) or key
diff --git a/skyflow/generated/rest/bin_lookup/__init__.py b/flowvault/skyflow_flowvault/generated/rest/flowservice/__init__.py
similarity index 100%
rename from skyflow/generated/rest/bin_lookup/__init__.py
rename to flowvault/skyflow_flowvault/generated/rest/flowservice/__init__.py
diff --git a/flowvault/skyflow_flowvault/generated/rest/flowservice/client.py b/flowvault/skyflow_flowvault/generated/rest/flowservice/client.py
new file mode 100644
index 00000000..321f69bd
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/flowservice/client.py
@@ -0,0 +1,855 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
+from ..core.request_options import RequestOptions
+from ..types.flow_enum_update_type import FlowEnumUpdateType
+from ..types.v_1_column_redactions import V1ColumnRedactions
+from ..types.v_1_delete_response import V1DeleteResponse
+from ..types.v_1_flow_delete_token_response import V1FlowDeleteTokenResponse
+from ..types.v_1_flow_detokenize_response import V1FlowDetokenizeResponse
+from ..types.v_1_flow_tokenize_request_object import V1FlowTokenizeRequestObject
+from ..types.v_1_flow_tokenize_response import V1FlowTokenizeResponse
+from ..types.v_1_flow_vault_metrics_response import V1FlowVaultMetricsResponse
+from ..types.v_1_get_request_data import V1GetRequestData
+from ..types.v_1_get_response import V1GetResponse
+from ..types.v_1_insert_record_data import V1InsertRecordData
+from ..types.v_1_insert_response import V1InsertResponse
+from ..types.v_1_token_group_redactions import V1TokenGroupRedactions
+from ..types.v_1_unique_value import V1UniqueValue
+from ..types.v_1_update_record_data import V1UpdateRecordData
+from ..types.v_1_update_response import V1UpdateResponse
+from ..types.v_1_upsert import V1Upsert
+from .raw_client import AsyncRawFlowserviceClient, RawFlowserviceClient
+
+# this is used as the default value for optional parameters
+OMIT = typing.cast(typing.Any, ...)
+
+
+class FlowserviceClient:
+ def __init__(self, *, client_wrapper: SyncClientWrapper):
+ self._raw_client = RawFlowserviceClient(client_wrapper=client_wrapper)
+
+ @property
+ def with_raw_response(self) -> RawFlowserviceClient:
+ """
+ Retrieves a raw implementation of this client that returns raw responses.
+
+ Returns
+ -------
+ RawFlowserviceClient
+ """
+ return self._raw_client
+
+ def delete(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ skyflow_i_ds: typing.Optional[typing.Sequence[str]] = OMIT,
+ unique_values: typing.Optional[typing.Sequence[V1UniqueValue]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1DeleteResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being deleted
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being deleted
+
+ skyflow_i_ds : typing.Optional[typing.Sequence[str]]
+ Skyflow ID for the record to be deleted
+
+ unique_values : typing.Optional[typing.Sequence[V1UniqueValue]]
+ List of unique constraint values to query records by data
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1DeleteResponse
+ A successful response.
+
+ Examples
+ --------
+ from skyflow import SkyflowAuth
+
+ client = SkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+ client.flowservice.delete()
+ """
+ _response = self._raw_client.delete(
+ vault_id=vault_id,
+ table_name=table_name,
+ skyflow_i_ds=skyflow_i_ds,
+ unique_values=unique_values,
+ request_options=request_options,
+ )
+ return _response.data
+
+ def get(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ skyflow_i_ds: typing.Optional[typing.Sequence[str]] = OMIT,
+ column_redactions: typing.Optional[typing.Sequence[V1ColumnRedactions]] = OMIT,
+ columns: typing.Optional[typing.Sequence[str]] = OMIT,
+ limit: typing.Optional[int] = OMIT,
+ offset: typing.Optional[int] = OMIT,
+ unique_values: typing.Optional[typing.Sequence[V1UniqueValue]] = OMIT,
+ records: typing.Optional[typing.Sequence[V1GetRequestData]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1GetResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being fetched
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being fetched
+
+ skyflow_i_ds : typing.Optional[typing.Sequence[str]]
+ Skyflow ID for the record to be fetched
+
+ column_redactions : typing.Optional[typing.Sequence[V1ColumnRedactions]]
+ List of columns to be redacted.
+
+ columns : typing.Optional[typing.Sequence[str]]
+ List of columns to be fetched.
+
+ limit : typing.Optional[int]
+ Limit for the number of records to be fetched
+
+ offset : typing.Optional[int]
+ Offset for the number of records to be fetched
+
+ unique_values : typing.Optional[typing.Sequence[V1UniqueValue]]
+ List of unique constraint values to query records by data
+
+ records : typing.Optional[typing.Sequence[V1GetRequestData]]
+ List of records to be fetched. This field contains tableName and skyflowIDs belonging to the table.
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1GetResponse
+ A successful response.
+
+ Examples
+ --------
+ from skyflow import SkyflowAuth
+
+ client = SkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+ client.flowservice.get()
+ """
+ _response = self._raw_client.get(
+ vault_id=vault_id,
+ table_name=table_name,
+ skyflow_i_ds=skyflow_i_ds,
+ column_redactions=column_redactions,
+ columns=columns,
+ limit=limit,
+ offset=offset,
+ unique_values=unique_values,
+ records=records,
+ request_options=request_options,
+ )
+ return _response.data
+
+ def insert(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ records: typing.Optional[typing.Sequence[V1InsertRecordData]] = OMIT,
+ upsert: typing.Optional[V1Upsert] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1InsertResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being inserted
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being inserted
+
+ records : typing.Optional[typing.Sequence[V1InsertRecordData]]
+ List of data row wise that is to be inserted in the vault
+
+ upsert : typing.Optional[V1Upsert]
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1InsertResponse
+ A successful response.
+
+ Examples
+ --------
+ from skyflow import SkyflowAuth
+
+ client = SkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+ client.flowservice.insert()
+ """
+ _response = self._raw_client.insert(
+ vault_id=vault_id, table_name=table_name, records=records, upsert=upsert, request_options=request_options
+ )
+ return _response.data
+
+ def update(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ records: typing.Optional[typing.Sequence[V1UpdateRecordData]] = OMIT,
+ update_type: typing.Optional[FlowEnumUpdateType] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1UpdateResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being updated
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being updated
+
+ records : typing.Optional[typing.Sequence[V1UpdateRecordData]]
+ List of data row wise that is to be updated in the vault
+
+ update_type : typing.Optional[FlowEnumUpdateType]
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1UpdateResponse
+ A successful response.
+
+ Examples
+ --------
+ from skyflow import SkyflowAuth
+
+ client = SkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+ client.flowservice.update()
+ """
+ _response = self._raw_client.update(
+ vault_id=vault_id,
+ table_name=table_name,
+ records=records,
+ update_type=update_type,
+ request_options=request_options,
+ )
+ return _response.data
+
+ def deletetoken(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ tokens: typing.Optional[typing.Sequence[str]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1FlowDeleteTokenResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ Vault ID
+
+ tokens : typing.Optional[typing.Sequence[str]]
+ Token value
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1FlowDeleteTokenResponse
+ A successful response.
+
+ Examples
+ --------
+ from skyflow import SkyflowAuth
+
+ client = SkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+ client.flowservice.deletetoken()
+ """
+ _response = self._raw_client.deletetoken(vault_id=vault_id, tokens=tokens, request_options=request_options)
+ return _response.data
+
+ def detokenize(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ tokens: typing.Optional[typing.Sequence[str]] = OMIT,
+ token_group_redactions: typing.Optional[typing.Sequence[V1TokenGroupRedactions]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1FlowDetokenizeResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where detokenizing
+
+ tokens : typing.Optional[typing.Sequence[str]]
+ Token to be detokenized
+
+ token_group_redactions : typing.Optional[typing.Sequence[V1TokenGroupRedactions]]
+ List of token groups to be redacted.
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1FlowDetokenizeResponse
+ A successful response.
+
+ Examples
+ --------
+ from skyflow import SkyflowAuth
+
+ client = SkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+ client.flowservice.detokenize()
+ """
+ _response = self._raw_client.detokenize(
+ vault_id=vault_id,
+ tokens=tokens,
+ token_group_redactions=token_group_redactions,
+ request_options=request_options,
+ )
+ return _response.data
+
+ def tokenize(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ data: typing.Optional[typing.Sequence[V1FlowTokenizeRequestObject]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1FlowTokenizeResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ Vault ID.
+
+ data : typing.Optional[typing.Sequence[V1FlowTokenizeRequestObject]]
+ Data to be tokenized
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1FlowTokenizeResponse
+ A successful response.
+
+ Examples
+ --------
+ from skyflow import SkyflowAuth
+
+ client = SkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+ client.flowservice.tokenize()
+ """
+ _response = self._raw_client.tokenize(vault_id=vault_id, data=data, request_options=request_options)
+ return _response.data
+
+ def flowvaultmetrics(
+ self, *, vault_id: typing.Optional[str] = OMIT, request_options: typing.Optional[RequestOptions] = None
+ ) -> V1FlowVaultMetricsResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault to get metrics for
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1FlowVaultMetricsResponse
+ A successful response.
+
+ Examples
+ --------
+ from skyflow import SkyflowAuth
+
+ client = SkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+ client.flowservice.flowvaultmetrics()
+ """
+ _response = self._raw_client.flowvaultmetrics(vault_id=vault_id, request_options=request_options)
+ return _response.data
+
+
+class AsyncFlowserviceClient:
+ def __init__(self, *, client_wrapper: AsyncClientWrapper):
+ self._raw_client = AsyncRawFlowserviceClient(client_wrapper=client_wrapper)
+
+ @property
+ def with_raw_response(self) -> AsyncRawFlowserviceClient:
+ """
+ Retrieves a raw implementation of this client that returns raw responses.
+
+ Returns
+ -------
+ AsyncRawFlowserviceClient
+ """
+ return self._raw_client
+
+ async def delete(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ skyflow_i_ds: typing.Optional[typing.Sequence[str]] = OMIT,
+ unique_values: typing.Optional[typing.Sequence[V1UniqueValue]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1DeleteResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being deleted
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being deleted
+
+ skyflow_i_ds : typing.Optional[typing.Sequence[str]]
+ Skyflow ID for the record to be deleted
+
+ unique_values : typing.Optional[typing.Sequence[V1UniqueValue]]
+ List of unique constraint values to query records by data
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1DeleteResponse
+ A successful response.
+
+ Examples
+ --------
+ import asyncio
+
+ from skyflow import AsyncSkyflowAuth
+
+ client = AsyncSkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+
+
+ async def main() -> None:
+ await client.flowservice.delete()
+
+
+ asyncio.run(main())
+ """
+ _response = await self._raw_client.delete(
+ vault_id=vault_id,
+ table_name=table_name,
+ skyflow_i_ds=skyflow_i_ds,
+ unique_values=unique_values,
+ request_options=request_options,
+ )
+ return _response.data
+
+ async def get(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ skyflow_i_ds: typing.Optional[typing.Sequence[str]] = OMIT,
+ column_redactions: typing.Optional[typing.Sequence[V1ColumnRedactions]] = OMIT,
+ columns: typing.Optional[typing.Sequence[str]] = OMIT,
+ limit: typing.Optional[int] = OMIT,
+ offset: typing.Optional[int] = OMIT,
+ unique_values: typing.Optional[typing.Sequence[V1UniqueValue]] = OMIT,
+ records: typing.Optional[typing.Sequence[V1GetRequestData]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1GetResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being fetched
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being fetched
+
+ skyflow_i_ds : typing.Optional[typing.Sequence[str]]
+ Skyflow ID for the record to be fetched
+
+ column_redactions : typing.Optional[typing.Sequence[V1ColumnRedactions]]
+ List of columns to be redacted.
+
+ columns : typing.Optional[typing.Sequence[str]]
+ List of columns to be fetched.
+
+ limit : typing.Optional[int]
+ Limit for the number of records to be fetched
+
+ offset : typing.Optional[int]
+ Offset for the number of records to be fetched
+
+ unique_values : typing.Optional[typing.Sequence[V1UniqueValue]]
+ List of unique constraint values to query records by data
+
+ records : typing.Optional[typing.Sequence[V1GetRequestData]]
+ List of records to be fetched. This field contains tableName and skyflowIDs belonging to the table.
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1GetResponse
+ A successful response.
+
+ Examples
+ --------
+ import asyncio
+
+ from skyflow import AsyncSkyflowAuth
+
+ client = AsyncSkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+
+
+ async def main() -> None:
+ await client.flowservice.get()
+
+
+ asyncio.run(main())
+ """
+ _response = await self._raw_client.get(
+ vault_id=vault_id,
+ table_name=table_name,
+ skyflow_i_ds=skyflow_i_ds,
+ column_redactions=column_redactions,
+ columns=columns,
+ limit=limit,
+ offset=offset,
+ unique_values=unique_values,
+ records=records,
+ request_options=request_options,
+ )
+ return _response.data
+
+ async def insert(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ records: typing.Optional[typing.Sequence[V1InsertRecordData]] = OMIT,
+ upsert: typing.Optional[V1Upsert] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1InsertResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being inserted
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being inserted
+
+ records : typing.Optional[typing.Sequence[V1InsertRecordData]]
+ List of data row wise that is to be inserted in the vault
+
+ upsert : typing.Optional[V1Upsert]
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1InsertResponse
+ A successful response.
+
+ Examples
+ --------
+ import asyncio
+
+ from skyflow import AsyncSkyflowAuth
+
+ client = AsyncSkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+
+
+ async def main() -> None:
+ await client.flowservice.insert()
+
+
+ asyncio.run(main())
+ """
+ _response = await self._raw_client.insert(
+ vault_id=vault_id, table_name=table_name, records=records, upsert=upsert, request_options=request_options
+ )
+ return _response.data
+
+ async def update(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ records: typing.Optional[typing.Sequence[V1UpdateRecordData]] = OMIT,
+ update_type: typing.Optional[FlowEnumUpdateType] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1UpdateResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being updated
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being updated
+
+ records : typing.Optional[typing.Sequence[V1UpdateRecordData]]
+ List of data row wise that is to be updated in the vault
+
+ update_type : typing.Optional[FlowEnumUpdateType]
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1UpdateResponse
+ A successful response.
+
+ Examples
+ --------
+ import asyncio
+
+ from skyflow import AsyncSkyflowAuth
+
+ client = AsyncSkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+
+
+ async def main() -> None:
+ await client.flowservice.update()
+
+
+ asyncio.run(main())
+ """
+ _response = await self._raw_client.update(
+ vault_id=vault_id,
+ table_name=table_name,
+ records=records,
+ update_type=update_type,
+ request_options=request_options,
+ )
+ return _response.data
+
+ async def deletetoken(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ tokens: typing.Optional[typing.Sequence[str]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1FlowDeleteTokenResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ Vault ID
+
+ tokens : typing.Optional[typing.Sequence[str]]
+ Token value
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1FlowDeleteTokenResponse
+ A successful response.
+
+ Examples
+ --------
+ import asyncio
+
+ from skyflow import AsyncSkyflowAuth
+
+ client = AsyncSkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+
+
+ async def main() -> None:
+ await client.flowservice.deletetoken()
+
+
+ asyncio.run(main())
+ """
+ _response = await self._raw_client.deletetoken(
+ vault_id=vault_id, tokens=tokens, request_options=request_options
+ )
+ return _response.data
+
+ async def detokenize(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ tokens: typing.Optional[typing.Sequence[str]] = OMIT,
+ token_group_redactions: typing.Optional[typing.Sequence[V1TokenGroupRedactions]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1FlowDetokenizeResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where detokenizing
+
+ tokens : typing.Optional[typing.Sequence[str]]
+ Token to be detokenized
+
+ token_group_redactions : typing.Optional[typing.Sequence[V1TokenGroupRedactions]]
+ List of token groups to be redacted.
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1FlowDetokenizeResponse
+ A successful response.
+
+ Examples
+ --------
+ import asyncio
+
+ from skyflow import AsyncSkyflowAuth
+
+ client = AsyncSkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+
+
+ async def main() -> None:
+ await client.flowservice.detokenize()
+
+
+ asyncio.run(main())
+ """
+ _response = await self._raw_client.detokenize(
+ vault_id=vault_id,
+ tokens=tokens,
+ token_group_redactions=token_group_redactions,
+ request_options=request_options,
+ )
+ return _response.data
+
+ async def tokenize(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ data: typing.Optional[typing.Sequence[V1FlowTokenizeRequestObject]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1FlowTokenizeResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ Vault ID.
+
+ data : typing.Optional[typing.Sequence[V1FlowTokenizeRequestObject]]
+ Data to be tokenized
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1FlowTokenizeResponse
+ A successful response.
+
+ Examples
+ --------
+ import asyncio
+
+ from skyflow import AsyncSkyflowAuth
+
+ client = AsyncSkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+
+
+ async def main() -> None:
+ await client.flowservice.tokenize()
+
+
+ asyncio.run(main())
+ """
+ _response = await self._raw_client.tokenize(vault_id=vault_id, data=data, request_options=request_options)
+ return _response.data
+
+ async def flowvaultmetrics(
+ self, *, vault_id: typing.Optional[str] = OMIT, request_options: typing.Optional[RequestOptions] = None
+ ) -> V1FlowVaultMetricsResponse:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault to get metrics for
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1FlowVaultMetricsResponse
+ A successful response.
+
+ Examples
+ --------
+ import asyncio
+
+ from skyflow import AsyncSkyflowAuth
+
+ client = AsyncSkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+
+
+ async def main() -> None:
+ await client.flowservice.flowvaultmetrics()
+
+
+ asyncio.run(main())
+ """
+ _response = await self._raw_client.flowvaultmetrics(vault_id=vault_id, request_options=request_options)
+ return _response.data
diff --git a/flowvault/skyflow_flowvault/generated/rest/flowservice/raw_client.py b/flowvault/skyflow_flowvault/generated/rest/flowservice/raw_client.py
new file mode 100644
index 00000000..7b005ea6
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/flowservice/raw_client.py
@@ -0,0 +1,1033 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+from json.decoder import JSONDecodeError
+
+from ..core.api_error import ApiError
+from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
+from ..core.http_response import AsyncHttpResponse, HttpResponse
+from ..core.pydantic_utilities import parse_obj_as
+from ..core.request_options import RequestOptions
+from ..core.serialization import convert_and_respect_annotation_metadata
+from ..types.flow_enum_update_type import FlowEnumUpdateType
+from ..types.v_1_column_redactions import V1ColumnRedactions
+from ..types.v_1_delete_response import V1DeleteResponse
+from ..types.v_1_flow_delete_token_response import V1FlowDeleteTokenResponse
+from ..types.v_1_flow_detokenize_response import V1FlowDetokenizeResponse
+from ..types.v_1_flow_tokenize_request_object import V1FlowTokenizeRequestObject
+from ..types.v_1_flow_tokenize_response import V1FlowTokenizeResponse
+from ..types.v_1_flow_vault_metrics_response import V1FlowVaultMetricsResponse
+from ..types.v_1_get_request_data import V1GetRequestData
+from ..types.v_1_get_response import V1GetResponse
+from ..types.v_1_insert_record_data import V1InsertRecordData
+from ..types.v_1_insert_response import V1InsertResponse
+from ..types.v_1_token_group_redactions import V1TokenGroupRedactions
+from ..types.v_1_unique_value import V1UniqueValue
+from ..types.v_1_update_record_data import V1UpdateRecordData
+from ..types.v_1_update_response import V1UpdateResponse
+from ..types.v_1_upsert import V1Upsert
+
+# this is used as the default value for optional parameters
+OMIT = typing.cast(typing.Any, ...)
+
+
+class RawFlowserviceClient:
+ def __init__(self, *, client_wrapper: SyncClientWrapper):
+ self._client_wrapper = client_wrapper
+
+ def delete(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ skyflow_i_ds: typing.Optional[typing.Sequence[str]] = OMIT,
+ unique_values: typing.Optional[typing.Sequence[V1UniqueValue]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> HttpResponse[V1DeleteResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being deleted
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being deleted
+
+ skyflow_i_ds : typing.Optional[typing.Sequence[str]]
+ Skyflow ID for the record to be deleted
+
+ unique_values : typing.Optional[typing.Sequence[V1UniqueValue]]
+ List of unique constraint values to query records by data
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ HttpResponse[V1DeleteResponse]
+ A successful response.
+ """
+ _response = self._client_wrapper.httpx_client.request(
+ "v2/records/delete",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "tableName": table_name,
+ "skyflowIDs": skyflow_i_ds,
+ "uniqueValues": convert_and_respect_annotation_metadata(
+ object_=unique_values, annotation=typing.Sequence[V1UniqueValue], direction="write"
+ ),
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1DeleteResponse,
+ parse_obj_as(
+ type_=V1DeleteResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return HttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ def get(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ skyflow_i_ds: typing.Optional[typing.Sequence[str]] = OMIT,
+ column_redactions: typing.Optional[typing.Sequence[V1ColumnRedactions]] = OMIT,
+ columns: typing.Optional[typing.Sequence[str]] = OMIT,
+ limit: typing.Optional[int] = OMIT,
+ offset: typing.Optional[int] = OMIT,
+ unique_values: typing.Optional[typing.Sequence[V1UniqueValue]] = OMIT,
+ records: typing.Optional[typing.Sequence[V1GetRequestData]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> HttpResponse[V1GetResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being fetched
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being fetched
+
+ skyflow_i_ds : typing.Optional[typing.Sequence[str]]
+ Skyflow ID for the record to be fetched
+
+ column_redactions : typing.Optional[typing.Sequence[V1ColumnRedactions]]
+ List of columns to be redacted.
+
+ columns : typing.Optional[typing.Sequence[str]]
+ List of columns to be fetched.
+
+ limit : typing.Optional[int]
+ Limit for the number of records to be fetched
+
+ offset : typing.Optional[int]
+ Offset for the number of records to be fetched
+
+ unique_values : typing.Optional[typing.Sequence[V1UniqueValue]]
+ List of unique constraint values to query records by data
+
+ records : typing.Optional[typing.Sequence[V1GetRequestData]]
+ List of records to be fetched. This field contains tableName and skyflowIDs belonging to the table.
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ HttpResponse[V1GetResponse]
+ A successful response.
+ """
+ _response = self._client_wrapper.httpx_client.request(
+ "v2/records/get",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "tableName": table_name,
+ "skyflowIDs": skyflow_i_ds,
+ "columnRedactions": convert_and_respect_annotation_metadata(
+ object_=column_redactions, annotation=typing.Sequence[V1ColumnRedactions], direction="write"
+ ),
+ "columns": columns,
+ "limit": limit,
+ "offset": offset,
+ "uniqueValues": convert_and_respect_annotation_metadata(
+ object_=unique_values, annotation=typing.Sequence[V1UniqueValue], direction="write"
+ ),
+ "records": convert_and_respect_annotation_metadata(
+ object_=records, annotation=typing.Sequence[V1GetRequestData], direction="write"
+ ),
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1GetResponse,
+ parse_obj_as(
+ type_=V1GetResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return HttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ def insert(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ records: typing.Optional[typing.Sequence[V1InsertRecordData]] = OMIT,
+ upsert: typing.Optional[V1Upsert] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> HttpResponse[V1InsertResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being inserted
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being inserted
+
+ records : typing.Optional[typing.Sequence[V1InsertRecordData]]
+ List of data row wise that is to be inserted in the vault
+
+ upsert : typing.Optional[V1Upsert]
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ HttpResponse[V1InsertResponse]
+ A successful response.
+ """
+ _response = self._client_wrapper.httpx_client.request(
+ "v2/records/insert",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "tableName": table_name,
+ "records": convert_and_respect_annotation_metadata(
+ object_=records, annotation=typing.Sequence[V1InsertRecordData], direction="write"
+ ),
+ "upsert": convert_and_respect_annotation_metadata(
+ object_=upsert, annotation=V1Upsert, direction="write"
+ ),
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1InsertResponse,
+ parse_obj_as(
+ type_=V1InsertResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return HttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ def update(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ records: typing.Optional[typing.Sequence[V1UpdateRecordData]] = OMIT,
+ update_type: typing.Optional[FlowEnumUpdateType] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> HttpResponse[V1UpdateResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being updated
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being updated
+
+ records : typing.Optional[typing.Sequence[V1UpdateRecordData]]
+ List of data row wise that is to be updated in the vault
+
+ update_type : typing.Optional[FlowEnumUpdateType]
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ HttpResponse[V1UpdateResponse]
+ A successful response.
+ """
+ _response = self._client_wrapper.httpx_client.request(
+ "v2/records/update",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "tableName": table_name,
+ "records": convert_and_respect_annotation_metadata(
+ object_=records, annotation=typing.Sequence[V1UpdateRecordData], direction="write"
+ ),
+ "updateType": update_type,
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1UpdateResponse,
+ parse_obj_as(
+ type_=V1UpdateResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return HttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ def deletetoken(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ tokens: typing.Optional[typing.Sequence[str]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> HttpResponse[V1FlowDeleteTokenResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ Vault ID
+
+ tokens : typing.Optional[typing.Sequence[str]]
+ Token value
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ HttpResponse[V1FlowDeleteTokenResponse]
+ A successful response.
+ """
+ _response = self._client_wrapper.httpx_client.request(
+ "v2/tokens/delete",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "tokens": tokens,
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1FlowDeleteTokenResponse,
+ parse_obj_as(
+ type_=V1FlowDeleteTokenResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return HttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ def detokenize(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ tokens: typing.Optional[typing.Sequence[str]] = OMIT,
+ token_group_redactions: typing.Optional[typing.Sequence[V1TokenGroupRedactions]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> HttpResponse[V1FlowDetokenizeResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where detokenizing
+
+ tokens : typing.Optional[typing.Sequence[str]]
+ Token to be detokenized
+
+ token_group_redactions : typing.Optional[typing.Sequence[V1TokenGroupRedactions]]
+ List of token groups to be redacted.
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ HttpResponse[V1FlowDetokenizeResponse]
+ A successful response.
+ """
+ _response = self._client_wrapper.httpx_client.request(
+ "v2/tokens/detokenize",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "tokens": tokens,
+ "tokenGroupRedactions": convert_and_respect_annotation_metadata(
+ object_=token_group_redactions,
+ annotation=typing.Sequence[V1TokenGroupRedactions],
+ direction="write",
+ ),
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1FlowDetokenizeResponse,
+ parse_obj_as(
+ type_=V1FlowDetokenizeResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return HttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ def tokenize(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ data: typing.Optional[typing.Sequence[V1FlowTokenizeRequestObject]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> HttpResponse[V1FlowTokenizeResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ Vault ID.
+
+ data : typing.Optional[typing.Sequence[V1FlowTokenizeRequestObject]]
+ Data to be tokenized
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ HttpResponse[V1FlowTokenizeResponse]
+ A successful response.
+ """
+ _response = self._client_wrapper.httpx_client.request(
+ "v2/tokens/tokenize",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "data": convert_and_respect_annotation_metadata(
+ object_=data, annotation=typing.Sequence[V1FlowTokenizeRequestObject], direction="write"
+ ),
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1FlowTokenizeResponse,
+ parse_obj_as(
+ type_=V1FlowTokenizeResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return HttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ def flowvaultmetrics(
+ self, *, vault_id: typing.Optional[str] = OMIT, request_options: typing.Optional[RequestOptions] = None
+ ) -> HttpResponse[V1FlowVaultMetricsResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault to get metrics for
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ HttpResponse[V1FlowVaultMetricsResponse]
+ A successful response.
+ """
+ _response = self._client_wrapper.httpx_client.request(
+ "v2/vaults/metrics",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1FlowVaultMetricsResponse,
+ parse_obj_as(
+ type_=V1FlowVaultMetricsResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return HttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+
+class AsyncRawFlowserviceClient:
+ def __init__(self, *, client_wrapper: AsyncClientWrapper):
+ self._client_wrapper = client_wrapper
+
+ async def delete(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ skyflow_i_ds: typing.Optional[typing.Sequence[str]] = OMIT,
+ unique_values: typing.Optional[typing.Sequence[V1UniqueValue]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> AsyncHttpResponse[V1DeleteResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being deleted
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being deleted
+
+ skyflow_i_ds : typing.Optional[typing.Sequence[str]]
+ Skyflow ID for the record to be deleted
+
+ unique_values : typing.Optional[typing.Sequence[V1UniqueValue]]
+ List of unique constraint values to query records by data
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ AsyncHttpResponse[V1DeleteResponse]
+ A successful response.
+ """
+ _response = await self._client_wrapper.httpx_client.request(
+ "v2/records/delete",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "tableName": table_name,
+ "skyflowIDs": skyflow_i_ds,
+ "uniqueValues": convert_and_respect_annotation_metadata(
+ object_=unique_values, annotation=typing.Sequence[V1UniqueValue], direction="write"
+ ),
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1DeleteResponse,
+ parse_obj_as(
+ type_=V1DeleteResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return AsyncHttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ async def get(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ skyflow_i_ds: typing.Optional[typing.Sequence[str]] = OMIT,
+ column_redactions: typing.Optional[typing.Sequence[V1ColumnRedactions]] = OMIT,
+ columns: typing.Optional[typing.Sequence[str]] = OMIT,
+ limit: typing.Optional[int] = OMIT,
+ offset: typing.Optional[int] = OMIT,
+ unique_values: typing.Optional[typing.Sequence[V1UniqueValue]] = OMIT,
+ records: typing.Optional[typing.Sequence[V1GetRequestData]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> AsyncHttpResponse[V1GetResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being fetched
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being fetched
+
+ skyflow_i_ds : typing.Optional[typing.Sequence[str]]
+ Skyflow ID for the record to be fetched
+
+ column_redactions : typing.Optional[typing.Sequence[V1ColumnRedactions]]
+ List of columns to be redacted.
+
+ columns : typing.Optional[typing.Sequence[str]]
+ List of columns to be fetched.
+
+ limit : typing.Optional[int]
+ Limit for the number of records to be fetched
+
+ offset : typing.Optional[int]
+ Offset for the number of records to be fetched
+
+ unique_values : typing.Optional[typing.Sequence[V1UniqueValue]]
+ List of unique constraint values to query records by data
+
+ records : typing.Optional[typing.Sequence[V1GetRequestData]]
+ List of records to be fetched. This field contains tableName and skyflowIDs belonging to the table.
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ AsyncHttpResponse[V1GetResponse]
+ A successful response.
+ """
+ _response = await self._client_wrapper.httpx_client.request(
+ "v2/records/get",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "tableName": table_name,
+ "skyflowIDs": skyflow_i_ds,
+ "columnRedactions": convert_and_respect_annotation_metadata(
+ object_=column_redactions, annotation=typing.Sequence[V1ColumnRedactions], direction="write"
+ ),
+ "columns": columns,
+ "limit": limit,
+ "offset": offset,
+ "uniqueValues": convert_and_respect_annotation_metadata(
+ object_=unique_values, annotation=typing.Sequence[V1UniqueValue], direction="write"
+ ),
+ "records": convert_and_respect_annotation_metadata(
+ object_=records, annotation=typing.Sequence[V1GetRequestData], direction="write"
+ ),
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1GetResponse,
+ parse_obj_as(
+ type_=V1GetResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return AsyncHttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ async def insert(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ records: typing.Optional[typing.Sequence[V1InsertRecordData]] = OMIT,
+ upsert: typing.Optional[V1Upsert] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> AsyncHttpResponse[V1InsertResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being inserted
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being inserted
+
+ records : typing.Optional[typing.Sequence[V1InsertRecordData]]
+ List of data row wise that is to be inserted in the vault
+
+ upsert : typing.Optional[V1Upsert]
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ AsyncHttpResponse[V1InsertResponse]
+ A successful response.
+ """
+ _response = await self._client_wrapper.httpx_client.request(
+ "v2/records/insert",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "tableName": table_name,
+ "records": convert_and_respect_annotation_metadata(
+ object_=records, annotation=typing.Sequence[V1InsertRecordData], direction="write"
+ ),
+ "upsert": convert_and_respect_annotation_metadata(
+ object_=upsert, annotation=V1Upsert, direction="write"
+ ),
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1InsertResponse,
+ parse_obj_as(
+ type_=V1InsertResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return AsyncHttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ async def update(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ table_name: typing.Optional[str] = OMIT,
+ records: typing.Optional[typing.Sequence[V1UpdateRecordData]] = OMIT,
+ update_type: typing.Optional[FlowEnumUpdateType] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> AsyncHttpResponse[V1UpdateResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being updated
+
+ table_name : typing.Optional[str]
+ Name of the table where data is being updated
+
+ records : typing.Optional[typing.Sequence[V1UpdateRecordData]]
+ List of data row wise that is to be updated in the vault
+
+ update_type : typing.Optional[FlowEnumUpdateType]
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ AsyncHttpResponse[V1UpdateResponse]
+ A successful response.
+ """
+ _response = await self._client_wrapper.httpx_client.request(
+ "v2/records/update",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "tableName": table_name,
+ "records": convert_and_respect_annotation_metadata(
+ object_=records, annotation=typing.Sequence[V1UpdateRecordData], direction="write"
+ ),
+ "updateType": update_type,
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1UpdateResponse,
+ parse_obj_as(
+ type_=V1UpdateResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return AsyncHttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ async def deletetoken(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ tokens: typing.Optional[typing.Sequence[str]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> AsyncHttpResponse[V1FlowDeleteTokenResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ Vault ID
+
+ tokens : typing.Optional[typing.Sequence[str]]
+ Token value
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ AsyncHttpResponse[V1FlowDeleteTokenResponse]
+ A successful response.
+ """
+ _response = await self._client_wrapper.httpx_client.request(
+ "v2/tokens/delete",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "tokens": tokens,
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1FlowDeleteTokenResponse,
+ parse_obj_as(
+ type_=V1FlowDeleteTokenResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return AsyncHttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ async def detokenize(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ tokens: typing.Optional[typing.Sequence[str]] = OMIT,
+ token_group_redactions: typing.Optional[typing.Sequence[V1TokenGroupRedactions]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> AsyncHttpResponse[V1FlowDetokenizeResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where detokenizing
+
+ tokens : typing.Optional[typing.Sequence[str]]
+ Token to be detokenized
+
+ token_group_redactions : typing.Optional[typing.Sequence[V1TokenGroupRedactions]]
+ List of token groups to be redacted.
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ AsyncHttpResponse[V1FlowDetokenizeResponse]
+ A successful response.
+ """
+ _response = await self._client_wrapper.httpx_client.request(
+ "v2/tokens/detokenize",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "tokens": tokens,
+ "tokenGroupRedactions": convert_and_respect_annotation_metadata(
+ object_=token_group_redactions,
+ annotation=typing.Sequence[V1TokenGroupRedactions],
+ direction="write",
+ ),
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1FlowDetokenizeResponse,
+ parse_obj_as(
+ type_=V1FlowDetokenizeResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return AsyncHttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ async def tokenize(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ data: typing.Optional[typing.Sequence[V1FlowTokenizeRequestObject]] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> AsyncHttpResponse[V1FlowTokenizeResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ Vault ID.
+
+ data : typing.Optional[typing.Sequence[V1FlowTokenizeRequestObject]]
+ Data to be tokenized
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ AsyncHttpResponse[V1FlowTokenizeResponse]
+ A successful response.
+ """
+ _response = await self._client_wrapper.httpx_client.request(
+ "v2/tokens/tokenize",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "data": convert_and_respect_annotation_metadata(
+ object_=data, annotation=typing.Sequence[V1FlowTokenizeRequestObject], direction="write"
+ ),
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1FlowTokenizeResponse,
+ parse_obj_as(
+ type_=V1FlowTokenizeResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return AsyncHttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+ async def flowvaultmetrics(
+ self, *, vault_id: typing.Optional[str] = OMIT, request_options: typing.Optional[RequestOptions] = None
+ ) -> AsyncHttpResponse[V1FlowVaultMetricsResponse]:
+ """
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault to get metrics for
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ AsyncHttpResponse[V1FlowVaultMetricsResponse]
+ A successful response.
+ """
+ _response = await self._client_wrapper.httpx_client.request(
+ "v2/vaults/metrics",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1FlowVaultMetricsResponse,
+ parse_obj_as(
+ type_=V1FlowVaultMetricsResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return AsyncHttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
diff --git a/skyflow/py.typed b/flowvault/skyflow_flowvault/generated/rest/py.typed
similarity index 100%
rename from skyflow/py.typed
rename to flowvault/skyflow_flowvault/generated/rest/py.typed
diff --git a/skyflow/generated/rest/guardrails/__init__.py b/flowvault/skyflow_flowvault/generated/rest/records/__init__.py
similarity index 100%
rename from skyflow/generated/rest/guardrails/__init__.py
rename to flowvault/skyflow_flowvault/generated/rest/records/__init__.py
diff --git a/flowvault/skyflow_flowvault/generated/rest/records/client.py b/flowvault/skyflow_flowvault/generated/rest/records/client.py
new file mode 100644
index 00000000..0503a99e
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/records/client.py
@@ -0,0 +1,131 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
+from ..core.request_options import RequestOptions
+from ..types.v_1_execute_query_response import V1ExecuteQueryResponse
+from .raw_client import AsyncRawRecordsClient, RawRecordsClient
+
+# this is used as the default value for optional parameters
+OMIT = typing.cast(typing.Any, ...)
+
+
+class RecordsClient:
+ def __init__(self, *, client_wrapper: SyncClientWrapper):
+ self._raw_client = RawRecordsClient(client_wrapper=client_wrapper)
+
+ @property
+ def with_raw_response(self) -> RawRecordsClient:
+ """
+ Retrieves a raw implementation of this client that returns raw responses.
+
+ Returns
+ -------
+ RawRecordsClient
+ """
+ return self._raw_client
+
+ def flow_service_execute_query(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ query: typing.Optional[str] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1ExecuteQueryResponse:
+ """
+ Executes a query on the specified vault.
+
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being inserted
+
+ query : typing.Optional[str]
+ Query to execute.
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1ExecuteQueryResponse
+ A successful response.
+
+ Examples
+ --------
+ from skyflow import SkyflowAuth
+
+ client = SkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+ client.records.flow_service_execute_query()
+ """
+ _response = self._raw_client.flow_service_execute_query(
+ vault_id=vault_id, query=query, request_options=request_options
+ )
+ return _response.data
+
+
+class AsyncRecordsClient:
+ def __init__(self, *, client_wrapper: AsyncClientWrapper):
+ self._raw_client = AsyncRawRecordsClient(client_wrapper=client_wrapper)
+
+ @property
+ def with_raw_response(self) -> AsyncRawRecordsClient:
+ """
+ Retrieves a raw implementation of this client that returns raw responses.
+
+ Returns
+ -------
+ AsyncRawRecordsClient
+ """
+ return self._raw_client
+
+ async def flow_service_execute_query(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ query: typing.Optional[str] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> V1ExecuteQueryResponse:
+ """
+ Executes a query on the specified vault.
+
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being inserted
+
+ query : typing.Optional[str]
+ Query to execute.
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ V1ExecuteQueryResponse
+ A successful response.
+
+ Examples
+ --------
+ import asyncio
+
+ from skyflow import AsyncSkyflowAuth
+
+ client = AsyncSkyflowAuth(
+ base_url="https://yourhost.com/path/to/api",
+ )
+
+
+ async def main() -> None:
+ await client.records.flow_service_execute_query()
+
+
+ asyncio.run(main())
+ """
+ _response = await self._raw_client.flow_service_execute_query(
+ vault_id=vault_id, query=query, request_options=request_options
+ )
+ return _response.data
diff --git a/flowvault/skyflow_flowvault/generated/rest/records/raw_client.py b/flowvault/skyflow_flowvault/generated/rest/records/raw_client.py
new file mode 100644
index 00000000..98a1365a
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/records/raw_client.py
@@ -0,0 +1,132 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+from json.decoder import JSONDecodeError
+
+from ..core.api_error import ApiError
+from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
+from ..core.http_response import AsyncHttpResponse, HttpResponse
+from ..core.pydantic_utilities import parse_obj_as
+from ..core.request_options import RequestOptions
+from ..types.v_1_execute_query_response import V1ExecuteQueryResponse
+
+# this is used as the default value for optional parameters
+OMIT = typing.cast(typing.Any, ...)
+
+
+class RawRecordsClient:
+ def __init__(self, *, client_wrapper: SyncClientWrapper):
+ self._client_wrapper = client_wrapper
+
+ def flow_service_execute_query(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ query: typing.Optional[str] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> HttpResponse[V1ExecuteQueryResponse]:
+ """
+ Executes a query on the specified vault.
+
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being inserted
+
+ query : typing.Optional[str]
+ Query to execute.
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ HttpResponse[V1ExecuteQueryResponse]
+ A successful response.
+ """
+ _response = self._client_wrapper.httpx_client.request(
+ "v2/query",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "query": query,
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1ExecuteQueryResponse,
+ parse_obj_as(
+ type_=V1ExecuteQueryResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return HttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+
+class AsyncRawRecordsClient:
+ def __init__(self, *, client_wrapper: AsyncClientWrapper):
+ self._client_wrapper = client_wrapper
+
+ async def flow_service_execute_query(
+ self,
+ *,
+ vault_id: typing.Optional[str] = OMIT,
+ query: typing.Optional[str] = OMIT,
+ request_options: typing.Optional[RequestOptions] = None,
+ ) -> AsyncHttpResponse[V1ExecuteQueryResponse]:
+ """
+ Executes a query on the specified vault.
+
+ Parameters
+ ----------
+ vault_id : typing.Optional[str]
+ ID of the vault where data is being inserted
+
+ query : typing.Optional[str]
+ Query to execute.
+
+ request_options : typing.Optional[RequestOptions]
+ Request-specific configuration.
+
+ Returns
+ -------
+ AsyncHttpResponse[V1ExecuteQueryResponse]
+ A successful response.
+ """
+ _response = await self._client_wrapper.httpx_client.request(
+ "v2/query",
+ method="POST",
+ json={
+ "vaultID": vault_id,
+ "query": query,
+ },
+ headers={
+ "content-type": "application/json",
+ },
+ request_options=request_options,
+ omit=OMIT,
+ )
+ try:
+ if 200 <= _response.status_code < 300:
+ _data = typing.cast(
+ V1ExecuteQueryResponse,
+ parse_obj_as(
+ type_=V1ExecuteQueryResponse, # type: ignore
+ object_=_response.json(),
+ ),
+ )
+ return AsyncHttpResponse(response=_response, data=_data)
+ _response_json = _response.json()
+ except JSONDecodeError:
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
+ raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/__init__.py b/flowvault/skyflow_flowvault/generated/rest/types/__init__.py
new file mode 100644
index 00000000..b088dc4c
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/__init__.py
@@ -0,0 +1,67 @@
+# This file was auto-generated by Fern from our API Definition.
+
+# isort: skip_file
+
+from .flow_enum_update_type import FlowEnumUpdateType
+from .flow_tokenize_response_object_token import FlowTokenizeResponseObjectToken
+from .googleprotobuf_any import GoogleprotobufAny
+from .protobuf_null_value import ProtobufNullValue
+from .rpc_status import RpcStatus
+from .v_1_column_redactions import V1ColumnRedactions
+from .v_1_delete_response import V1DeleteResponse
+from .v_1_delete_response_object import V1DeleteResponseObject
+from .v_1_delete_token_response_object import V1DeleteTokenResponseObject
+from .v_1_execute_query_record_response import V1ExecuteQueryRecordResponse
+from .v_1_execute_query_response import V1ExecuteQueryResponse
+from .v_1_execute_query_response_metadata import V1ExecuteQueryResponseMetadata
+from .v_1_flow_delete_token_response import V1FlowDeleteTokenResponse
+from .v_1_flow_detokenize_response import V1FlowDetokenizeResponse
+from .v_1_flow_detokenize_response_object import V1FlowDetokenizeResponseObject
+from .v_1_flow_tokenize_request_object import V1FlowTokenizeRequestObject
+from .v_1_flow_tokenize_response import V1FlowTokenizeResponse
+from .v_1_flow_tokenize_response_object import V1FlowTokenizeResponseObject
+from .v_1_flow_vault_metrics_data import V1FlowVaultMetricsData
+from .v_1_flow_vault_metrics_response import V1FlowVaultMetricsResponse
+from .v_1_get_request_data import V1GetRequestData
+from .v_1_get_response import V1GetResponse
+from .v_1_insert_record_data import V1InsertRecordData
+from .v_1_insert_response import V1InsertResponse
+from .v_1_record_response_object import V1RecordResponseObject
+from .v_1_token_group_redactions import V1TokenGroupRedactions
+from .v_1_unique_value import V1UniqueValue
+from .v_1_update_record_data import V1UpdateRecordData
+from .v_1_update_response import V1UpdateResponse
+from .v_1_upsert import V1Upsert
+
+__all__ = [
+ "FlowEnumUpdateType",
+ "FlowTokenizeResponseObjectToken",
+ "GoogleprotobufAny",
+ "ProtobufNullValue",
+ "RpcStatus",
+ "V1ColumnRedactions",
+ "V1DeleteResponse",
+ "V1DeleteResponseObject",
+ "V1DeleteTokenResponseObject",
+ "V1ExecuteQueryRecordResponse",
+ "V1ExecuteQueryResponse",
+ "V1ExecuteQueryResponseMetadata",
+ "V1FlowDeleteTokenResponse",
+ "V1FlowDetokenizeResponse",
+ "V1FlowDetokenizeResponseObject",
+ "V1FlowTokenizeRequestObject",
+ "V1FlowTokenizeResponse",
+ "V1FlowTokenizeResponseObject",
+ "V1FlowVaultMetricsData",
+ "V1FlowVaultMetricsResponse",
+ "V1GetRequestData",
+ "V1GetResponse",
+ "V1InsertRecordData",
+ "V1InsertResponse",
+ "V1RecordResponseObject",
+ "V1TokenGroupRedactions",
+ "V1UniqueValue",
+ "V1UpdateRecordData",
+ "V1UpdateResponse",
+ "V1Upsert",
+]
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/flow_enum_update_type.py b/flowvault/skyflow_flowvault/generated/rest/types/flow_enum_update_type.py
new file mode 100644
index 00000000..01b2bab9
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/flow_enum_update_type.py
@@ -0,0 +1,5 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+FlowEnumUpdateType = typing.Union[typing.Literal["UPDATE", "REPLACE"], typing.Any]
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/flow_tokenize_response_object_token.py b/flowvault/skyflow_flowvault/generated/rest/types/flow_tokenize_response_object_token.py
new file mode 100644
index 00000000..928a6606
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/flow_tokenize_response_object_token.py
@@ -0,0 +1,43 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class FlowTokenizeResponseObjectToken(UniversalBaseModel):
+ token_group_name: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="tokenGroupName")] = (
+ pydantic.Field(default=None)
+ )
+ """
+ Token group Name
+ """
+
+ token: typing.Optional[str] = pydantic.Field(default=None)
+ """
+ Token value
+ """
+
+ error: typing.Optional[str] = pydantic.Field(default=None)
+ """
+ Error if tokenization failed
+ """
+
+ http_code: typing_extensions.Annotated[typing.Optional[int], FieldMetadata(alias="httpCode")] = pydantic.Field(
+ default=None
+ )
+ """
+ HTTP status code of the response
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/googleprotobuf_any.py b/flowvault/skyflow_flowvault/generated/rest/types/googleprotobuf_any.py
new file mode 100644
index 00000000..aebcc5b9
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/googleprotobuf_any.py
@@ -0,0 +1,139 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class GoogleprotobufAny(UniversalBaseModel):
+ """
+ `Any` contains an arbitrary serialized protocol buffer message along with a
+ URL that describes the type of the serialized message.
+
+ Protobuf library provides support to pack/unpack Any values in the form
+ of utility functions or additional generated methods of the Any type.
+
+ Example 1: Pack and unpack a message in C++.
+
+ Foo foo = ...;
+ Any any;
+ any.PackFrom(foo);
+ ...
+ if (any.UnpackTo(&foo)) {
+ ...
+ }
+
+ Example 2: Pack and unpack a message in Java.
+
+ Foo foo = ...;
+ Any any = Any.pack(foo);
+ ...
+ if (any.is(Foo.class)) {
+ foo = any.unpack(Foo.class);
+ }
+ // or ...
+ if (any.isSameTypeAs(Foo.getDefaultInstance())) {
+ foo = any.unpack(Foo.getDefaultInstance());
+ }
+
+ Example 3: Pack and unpack a message in Python.
+
+ foo = Foo(...)
+ any = Any()
+ any.Pack(foo)
+ ...
+ if any.Is(Foo.DESCRIPTOR):
+ any.Unpack(foo)
+ ...
+
+ Example 4: Pack and unpack a message in Go
+
+ foo := &pb.Foo{...}
+ any, err := anypb.New(foo)
+ if err != nil {
+ ...
+ }
+ ...
+ foo := &pb.Foo{}
+ if err := any.UnmarshalTo(foo); err != nil {
+ ...
+ }
+
+ The pack methods provided by protobuf library will by default use
+ 'type.googleapis.com/full.type.name' as the type URL and the unpack
+ methods only use the fully qualified type name after the last '/'
+ in the type URL, for example "foo.bar.com/x/y.z" will yield type
+ name "y.z".
+
+ JSON
+ ====
+ The JSON representation of an `Any` value uses the regular
+ representation of the deserialized, embedded message, with an
+ additional field `@type` which contains the type URL. Example:
+
+ package google.profile;
+ message Person {
+ string first_name = 1;
+ string last_name = 2;
+ }
+
+ {
+ "@type": "type.googleapis.com/google.profile.Person",
+ "firstName": ,
+ "lastName":
+ }
+
+ If the embedded message type is well-known and has a custom JSON
+ representation, that representation will be embedded adding a field
+ `value` which holds the custom JSON in addition to the `@type`
+ field. Example (for message [google.protobuf.Duration][]):
+
+ {
+ "@type": "type.googleapis.com/google.protobuf.Duration",
+ "value": "1.212s"
+ }
+ """
+
+ type: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="@type")] = pydantic.Field(default=None)
+ """
+ A URL/resource name that uniquely identifies the type of the serialized
+ protocol buffer message. This string must contain at least
+ one "/" character. The last segment of the URL's path must represent
+ the fully qualified name of the type (as in
+ `path/google.protobuf.Duration`). The name should be in a canonical form
+ (e.g., leading "." is not accepted).
+
+ In practice, teams usually precompile into the binary all types that they
+ expect it to use in the context of Any. However, for URLs which use the
+ scheme `http`, `https`, or no scheme, one can optionally set up a type
+ server that maps type URLs to message definitions as follows:
+
+ * If no scheme is provided, `https` is assumed.
+ * An HTTP GET on the URL must yield a [google.protobuf.Type][]
+ value in binary format, or produce an error.
+ * Applications are allowed to cache lookup results based on the
+ URL, or have them precompiled into a binary to avoid any
+ lookup. Therefore, binary compatibility needs to be preserved
+ on changes to types. (Use versioned type names to manage
+ breaking changes.)
+
+ Note: this functionality is not currently available in the official
+ protobuf release, and it is not used for type URLs beginning with
+ type.googleapis.com. As of May 2023, there are no widely used type server
+ implementations and no plans to implement one.
+
+ Schemes other than `http`, `https` (or the empty scheme) might be
+ used with implementation specific semantics.
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/protobuf_null_value.py b/flowvault/skyflow_flowvault/generated/rest/types/protobuf_null_value.py
new file mode 100644
index 00000000..7a4d590f
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/protobuf_null_value.py
@@ -0,0 +1,5 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+ProtobufNullValue = typing.Literal["NULL_VALUE"]
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/rpc_status.py b/flowvault/skyflow_flowvault/generated/rest/types/rpc_status.py
new file mode 100644
index 00000000..cf324547
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/rpc_status.py
@@ -0,0 +1,22 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from .googleprotobuf_any import GoogleprotobufAny
+
+
+class RpcStatus(UniversalBaseModel):
+ code: typing.Optional[int] = None
+ message: typing.Optional[str] = None
+ details: typing.Optional[typing.List[GoogleprotobufAny]] = None
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_column_redactions.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_column_redactions.py
new file mode 100644
index 00000000..65d089a5
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_column_redactions.py
@@ -0,0 +1,31 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class V1ColumnRedactions(UniversalBaseModel):
+ column_name: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="columnName")] = pydantic.Field(
+ default=None
+ )
+ """
+ Name of the column to be redacted
+ """
+
+ redaction: typing.Optional[str] = pydantic.Field(default=None)
+ """
+ Name of the redaction. Eg: `plain_text`, `redacted`, `mask1`
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_delete_response.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_delete_response.py
new file mode 100644
index 00000000..9b281978
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_delete_response.py
@@ -0,0 +1,23 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from .v_1_delete_response_object import V1DeleteResponseObject
+
+
+class V1DeleteResponse(UniversalBaseModel):
+ records: typing.Optional[typing.List[V1DeleteResponseObject]] = pydantic.Field(default=None)
+ """
+ List of deleted records with skyflow ID and any partial errors.
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_delete_response_object.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_delete_response_object.py
new file mode 100644
index 00000000..eda4c5ab
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_delete_response_object.py
@@ -0,0 +1,38 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class V1DeleteResponseObject(UniversalBaseModel):
+ skyflow_id: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="skyflowID")] = pydantic.Field(
+ default=None
+ )
+ """
+ Skyflow ID for the deleted record
+ """
+
+ error: typing.Optional[str] = pydantic.Field(default=None)
+ """
+ Partial Error message if any
+ """
+
+ http_code: typing_extensions.Annotated[typing.Optional[int], FieldMetadata(alias="httpCode")] = pydantic.Field(
+ default=None
+ )
+ """
+ HTTP status code of the response
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_delete_token_response_object.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_delete_token_response_object.py
new file mode 100644
index 00000000..2a482ec0
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_delete_token_response_object.py
@@ -0,0 +1,36 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class V1DeleteTokenResponseObject(UniversalBaseModel):
+ value: typing.Optional[str] = pydantic.Field(default=None)
+ """
+ Token value
+ """
+
+ error: typing.Optional[str] = pydantic.Field(default=None)
+ """
+ Error if deletion failed
+ """
+
+ http_code: typing_extensions.Annotated[typing.Optional[int], FieldMetadata(alias="httpCode")] = pydantic.Field(
+ default=None
+ )
+ """
+ HTTP status code of the response
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_execute_query_record_response.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_execute_query_record_response.py
new file mode 100644
index 00000000..30de3867
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_execute_query_record_response.py
@@ -0,0 +1,22 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+
+
+class V1ExecuteQueryRecordResponse(UniversalBaseModel):
+ data: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
+ """
+ Fields and values for the record. For example, `{'field_1':'value_1', 'field_2':'value_2'}`.
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_execute_query_response.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_execute_query_response.py
new file mode 100644
index 00000000..17caef33
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_execute_query_response.py
@@ -0,0 +1,26 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from .v_1_execute_query_record_response import V1ExecuteQueryRecordResponse
+from .v_1_execute_query_response_metadata import V1ExecuteQueryResponseMetadata
+
+
+class V1ExecuteQueryResponse(UniversalBaseModel):
+ records: typing.Optional[typing.List[V1ExecuteQueryRecordResponse]] = pydantic.Field(default=None)
+ """
+ Records corresponding to the specified query.
+ """
+
+ metadata: typing.Optional[V1ExecuteQueryResponseMetadata] = None
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_execute_query_response_metadata.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_execute_query_response_metadata.py
new file mode 100644
index 00000000..3eb0e86c
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_execute_query_response_metadata.py
@@ -0,0 +1,22 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+
+
+class V1ExecuteQueryResponseMetadata(UniversalBaseModel):
+ columns: typing.Optional[typing.List[str]] = pydantic.Field(default=None)
+ """
+ Return columns for the query
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_delete_token_response.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_delete_token_response.py
new file mode 100644
index 00000000..9129dbff
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_delete_token_response.py
@@ -0,0 +1,23 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from .v_1_delete_token_response_object import V1DeleteTokenResponseObject
+
+
+class V1FlowDeleteTokenResponse(UniversalBaseModel):
+ tokens: typing.Optional[typing.List[V1DeleteTokenResponseObject]] = pydantic.Field(default=None)
+ """
+ Tokens data for Delete
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_detokenize_response.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_detokenize_response.py
new file mode 100644
index 00000000..47ab50dd
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_detokenize_response.py
@@ -0,0 +1,23 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from .v_1_flow_detokenize_response_object import V1FlowDetokenizeResponseObject
+
+
+class V1FlowDetokenizeResponse(UniversalBaseModel):
+ response: typing.Optional[typing.List[V1FlowDetokenizeResponseObject]] = pydantic.Field(default=None)
+ """
+ Detokenized data
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_detokenize_response_object.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_detokenize_response_object.py
new file mode 100644
index 00000000..382a2b1a
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_detokenize_response_object.py
@@ -0,0 +1,53 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class V1FlowDetokenizeResponseObject(UniversalBaseModel):
+ token: typing.Optional[str] = pydantic.Field(default=None)
+ """
+ Token to be detokenized
+ """
+
+ value: typing.Optional[typing.Optional[typing.Any]] = pydantic.Field(default=None)
+ """
+ Detokenized value for the token
+ """
+
+ token_group_name: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="tokenGroupName")] = (
+ pydantic.Field(default=None)
+ )
+ """
+ Token group name
+ """
+
+ error: typing.Optional[str] = pydantic.Field(default=None)
+ """
+ Error if detokenization failed
+ """
+
+ http_code: typing_extensions.Annotated[typing.Optional[int], FieldMetadata(alias="httpCode")] = pydantic.Field(
+ default=None
+ )
+ """
+ HTTP status code of the response
+ """
+
+ metadata: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
+ """
+ Additional metadata associated with the token, such as tableName or skyflowID
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_tokenize_request_object.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_tokenize_request_object.py
new file mode 100644
index 00000000..42a926ee
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_tokenize_request_object.py
@@ -0,0 +1,36 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class V1FlowTokenizeRequestObject(UniversalBaseModel):
+ value: typing.Optional[typing.Optional[typing.Any]] = pydantic.Field(default=None)
+ """
+ Token Value
+ """
+
+ token_group_names: typing_extensions.Annotated[
+ typing.Optional[typing.List[str]], FieldMetadata(alias="tokenGroupNames")
+ ] = pydantic.Field(default=None)
+ """
+ List of token group names
+ """
+
+ token: typing.Optional[typing.Optional[typing.Any]] = pydantic.Field(default=None)
+ """
+ Token for the value, in case of BYOT.
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_tokenize_response.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_tokenize_response.py
new file mode 100644
index 00000000..88616410
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_tokenize_response.py
@@ -0,0 +1,23 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from .v_1_flow_tokenize_response_object import V1FlowTokenizeResponseObject
+
+
+class V1FlowTokenizeResponse(UniversalBaseModel):
+ response: typing.Optional[typing.List[V1FlowTokenizeResponseObject]] = pydantic.Field(default=None)
+ """
+ Tokenized data
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_tokenize_response_object.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_tokenize_response_object.py
new file mode 100644
index 00000000..e77e153b
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_tokenize_response_object.py
@@ -0,0 +1,28 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from .flow_tokenize_response_object_token import FlowTokenizeResponseObjectToken
+
+
+class V1FlowTokenizeResponseObject(UniversalBaseModel):
+ value: typing.Optional[typing.Optional[typing.Any]] = pydantic.Field(default=None)
+ """
+ Value of token
+ """
+
+ tokens: typing.Optional[typing.List[FlowTokenizeResponseObjectToken]] = pydantic.Field(default=None)
+ """
+ Token value
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_vault_metrics_data.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_vault_metrics_data.py
new file mode 100644
index 00000000..f8611c37
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_vault_metrics_data.py
@@ -0,0 +1,22 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+
+
+class V1FlowVaultMetricsData(UniversalBaseModel):
+ tables: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
+ """
+ Map of table names to their metrics
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_vault_metrics_response.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_vault_metrics_response.py
new file mode 100644
index 00000000..5234fd91
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_flow_vault_metrics_response.py
@@ -0,0 +1,24 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from .v_1_flow_vault_metrics_data import V1FlowVaultMetricsData
+
+
+class V1FlowVaultMetricsResponse(UniversalBaseModel):
+ data: typing.Optional[V1FlowVaultMetricsData] = None
+ error: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
+ """
+ Error information, if any
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_get_request_data.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_get_request_data.py
new file mode 100644
index 00000000..caf815b6
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_get_request_data.py
@@ -0,0 +1,54 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+from .v_1_column_redactions import V1ColumnRedactions
+from .v_1_unique_value import V1UniqueValue
+
+
+class V1GetRequestData(UniversalBaseModel):
+ table_name: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="tableName")] = pydantic.Field(
+ default=None
+ )
+ """
+ Name of the table where data is being fetched
+ """
+
+ skyflow_i_ds: typing_extensions.Annotated[typing.Optional[typing.List[str]], FieldMetadata(alias="skyflowIDs")] = (
+ pydantic.Field(default=None)
+ )
+ """
+ Skyflow ID for the record to be fetched
+ """
+
+ column_redactions: typing_extensions.Annotated[
+ typing.Optional[typing.List[V1ColumnRedactions]], FieldMetadata(alias="columnRedactions")
+ ] = pydantic.Field(default=None)
+ """
+ List of columns to be redacted.
+ """
+
+ columns: typing.Optional[typing.List[str]] = pydantic.Field(default=None)
+ """
+ List of columns to be fetched.
+ """
+
+ unique_values: typing_extensions.Annotated[
+ typing.Optional[typing.List[V1UniqueValue]], FieldMetadata(alias="uniqueValues")
+ ] = pydantic.Field(default=None)
+ """
+ List of unique constraint values to query records by data
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_get_response.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_get_response.py
new file mode 100644
index 00000000..ab966469
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_get_response.py
@@ -0,0 +1,23 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from .v_1_record_response_object import V1RecordResponseObject
+
+
+class V1GetResponse(UniversalBaseModel):
+ records: typing.Optional[typing.List[V1RecordResponseObject]] = pydantic.Field(default=None)
+ """
+ List of fetched records with skyflow ID, tokens, data, and any partial errors
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_insert_record_data.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_insert_record_data.py
new file mode 100644
index 00000000..063626d3
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_insert_record_data.py
@@ -0,0 +1,39 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+from .v_1_upsert import V1Upsert
+
+
+class V1InsertRecordData(UniversalBaseModel):
+ data: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
+ """
+ Columns names and values
+ """
+
+ tokens: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
+ """
+ undocumented_field; Tokens data for the columns if any
+ """
+
+ table_name: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="tableName")] = pydantic.Field(
+ default=None
+ )
+ """
+ Table name for the record
+ """
+
+ upsert: typing.Optional[V1Upsert] = None
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_insert_response.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_insert_response.py
new file mode 100644
index 00000000..bac58b52
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_insert_response.py
@@ -0,0 +1,23 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from .v_1_record_response_object import V1RecordResponseObject
+
+
+class V1InsertResponse(UniversalBaseModel):
+ records: typing.Optional[typing.List[V1RecordResponseObject]] = pydantic.Field(default=None)
+ """
+ List of inserted records with skyflow ID, tokens, data, and any partial errors.
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_record_response_object.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_record_response_object.py
new file mode 100644
index 00000000..0f02a93f
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_record_response_object.py
@@ -0,0 +1,62 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class V1RecordResponseObject(UniversalBaseModel):
+ skyflow_id: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="skyflowID")] = pydantic.Field(
+ default=None
+ )
+ """
+ Skyflow ID for the inserted record
+ """
+
+ tokens: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
+ """
+ Tokens data for the columns if any
+ """
+
+ data: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
+ """
+ Columns names and values
+ """
+
+ hashed_data: typing_extensions.Annotated[
+ typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]], FieldMetadata(alias="hashedData")
+ ] = pydantic.Field(default=None)
+ """
+ Hashed Data for the columns if any
+ """
+
+ error: typing.Optional[str] = pydantic.Field(default=None)
+ """
+ Partial Error message if any
+ """
+
+ http_code: typing_extensions.Annotated[typing.Optional[int], FieldMetadata(alias="httpCode")] = pydantic.Field(
+ default=None
+ )
+ """
+ HTTP status code of the response
+ """
+
+ table_name: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="tableName")] = pydantic.Field(
+ default=None
+ )
+ """
+ Name of the table record belongs to
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_token_group_redactions.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_token_group_redactions.py
new file mode 100644
index 00000000..69263a19
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_token_group_redactions.py
@@ -0,0 +1,31 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class V1TokenGroupRedactions(UniversalBaseModel):
+ token_group_name: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="tokenGroupName")] = (
+ pydantic.Field(default=None)
+ )
+ """
+ Name of the token group to be redacted
+ """
+
+ redaction: typing.Optional[str] = pydantic.Field(default=None)
+ """
+ Name of the redaction. Eg: `plain_text`, `redacted`, `mask1`
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_unique_value.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_unique_value.py
new file mode 100644
index 00000000..e0cfa021
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_unique_value.py
@@ -0,0 +1,22 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+
+
+class V1UniqueValue(UniversalBaseModel):
+ data: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
+ """
+ Columns names and values for unique value entry
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_update_record_data.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_update_record_data.py
new file mode 100644
index 00000000..19622eab
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_update_record_data.py
@@ -0,0 +1,43 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class V1UpdateRecordData(UniversalBaseModel):
+ skyflow_id: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="skyflowID")] = pydantic.Field(
+ default=None
+ )
+ """
+ Skyflow ID for the record to be updated
+ """
+
+ data: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
+ """
+ List of data row wise that is to be updated in the vault
+ """
+
+ tokens: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
+ """
+ undocumented_field; Tokens data for the columns if any
+ """
+
+ table_name: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="tableName")] = pydantic.Field(
+ default=None
+ )
+ """
+ Table name for the record
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_update_response.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_update_response.py
new file mode 100644
index 00000000..4f4eb228
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_update_response.py
@@ -0,0 +1,23 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from .v_1_record_response_object import V1RecordResponseObject
+
+
+class V1UpdateResponse(UniversalBaseModel):
+ records: typing.Optional[typing.List[V1RecordResponseObject]] = pydantic.Field(default=None)
+ """
+ List of updated records with skyflow ID, tokens, data, and any partial errors
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/types/v_1_upsert.py b/flowvault/skyflow_flowvault/generated/rest/types/v_1_upsert.py
new file mode 100644
index 00000000..f9531a37
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/types/v_1_upsert.py
@@ -0,0 +1,30 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+from .flow_enum_update_type import FlowEnumUpdateType
+
+
+class V1Upsert(UniversalBaseModel):
+ update_type: typing_extensions.Annotated[typing.Optional[FlowEnumUpdateType], FieldMetadata(alias="updateType")] = (
+ None
+ )
+ unique_columns: typing_extensions.Annotated[
+ typing.Optional[typing.List[str]], FieldMetadata(alias="uniqueColumns")
+ ] = pydantic.Field(default=None)
+ """
+ Name of a unique columns in the table. Uses upsert operations to check if a record exists based on the unique column's value. If a matching record exists, the record updates with the values you provide. If a matching record doesn't exist, the upsert operation inserts a new record.
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/flowvault/skyflow_flowvault/generated/rest/version.py b/flowvault/skyflow_flowvault/generated/rest/version.py
new file mode 100644
index 00000000..82f62b47
--- /dev/null
+++ b/flowvault/skyflow_flowvault/generated/rest/version.py
@@ -0,0 +1,6 @@
+# NOTE: hand-patched, not Fern-generated content. Fern originally emitted a runtime
+# metadata.version("skyflow.generated.rest") lookup here, but this code is bundled into the
+# skyflow (v2) and v3 wheels via a build_py hook rather than published under that distribution
+# name, so the lookup always raised PackageNotFoundError on import. Hardcoded until the Fern
+# generator config (skyflow-fern-config) is updated to stop emitting a runtime lookup here.
+__version__ = "0.0.10"
diff --git a/flowvault/skyflow_flowvault/service_account/__init__.py b/flowvault/skyflow_flowvault/service_account/__init__.py
new file mode 100644
index 00000000..1563c2fe
--- /dev/null
+++ b/flowvault/skyflow_flowvault/service_account/__init__.py
@@ -0,0 +1,15 @@
+from common.service_account import (
+ generate_bearer_token,
+ generate_bearer_token_from_creds,
+ is_expired,
+ generate_signed_data_tokens,
+ generate_signed_data_tokens_from_creds,
+)
+
+__all__ = [
+ "generate_bearer_token",
+ "generate_bearer_token_from_creds",
+ "is_expired",
+ "generate_signed_data_tokens",
+ "generate_signed_data_tokens_from_creds",
+]
diff --git a/flowvault/skyflow_flowvault/utils/__init__.py b/flowvault/skyflow_flowvault/utils/__init__.py
new file mode 100644
index 00000000..3b008cc5
--- /dev/null
+++ b/flowvault/skyflow_flowvault/utils/__init__.py
@@ -0,0 +1,9 @@
+# Must be imported before anything that pulls in common.utils -- see its _skyflow_messages.py
+# comment for why the load order matters.
+from ._version import SDK_VERSION
+
+from common.utils import LogLevel, Env
+
+from .enums import UpsertType, EnvUrls
+from ._skyflow_messages import SkyflowMessages
+from ._utils import get_metrics, get_vault_url
diff --git a/flowvault/skyflow_flowvault/utils/_skyflow_messages.py b/flowvault/skyflow_flowvault/utils/_skyflow_messages.py
new file mode 100644
index 00000000..35596296
--- /dev/null
+++ b/flowvault/skyflow_flowvault/utils/_skyflow_messages.py
@@ -0,0 +1,57 @@
+from enum import Enum
+
+try:
+ from ._version import SDK_VERSION
+except ImportError: # pragma: no cover
+ SDK_VERSION = "0.0.0"
+
+error_prefix = f"Skyflow Python SDK {SDK_VERSION}"
+INFO = "INFO"
+WARN = "WARN"
+ERROR = "ERROR"
+
+
+class SkyflowMessages:
+ """v3's own operation-specific message catalog, mirroring v2's per-variant pattern. Generic
+ infrastructure text lives in common.utils.SkyflowMessages instead."""
+
+ class Error(Enum):
+ EMPTY_RECORDS_IN_INSERT = f"{error_prefix} Insert failed. Specify at least one record to insert."
+ INVALID_RECORDS_TYPE_IN_INSERT = f"{error_prefix} Insert failed. 'records' must be a list of dicts."
+ INVALID_RECORD_DATA_IN_INSERT = f"{error_prefix} Insert failed. Each record's 'values' must be a non-empty dict."
+ INVALID_TABLE_NAME_IN_INSERT = f"{error_prefix} Insert failed. 'table' must be a non-empty string."
+ INVALID_UPSERT_TYPE_IN_INSERT = f"{error_prefix} Insert failed. 'upsert' must be a dict."
+ INVALID_UPSERT_UNIQUE_COLUMNS_IN_INSERT = f"{error_prefix} Insert failed. Upsert's 'unique_columns' must be a non-empty list of strings."
+ INVALID_UPSERT_UPDATE_TYPE_IN_INSERT = f"{error_prefix} Insert failed. Upsert's 'update_type' must be an UpsertType value."
+ TOO_MANY_RECORDS_IN_INSERT = f"{error_prefix} Insert failed. A single insert request cannot contain more than 10000 records."
+ TABLE_NAME_IN_BOTH_PLACES_IN_INSERT = (
+ f"{error_prefix} Insert failed. 'table' cannot be set on InsertRequest at the same "
+ "time as any record's 'table' -- the vault accepts a table name outside the records "
+ "(request-level, applying to all of them) or inside each record, but not both at once."
+ )
+ TABLE_NAME_MISSING_IN_INSERT = (
+ f"{error_prefix} Insert failed. 'table' is not set on InsertRequest, so every record "
+ "must set its own 'table' -- either set 'table' once at the request level, or set it "
+ "individually on every record."
+ )
+ RECORD_LEVEL_UPSERT_NOT_ALLOWED_IN_INSERT = (
+ f"{error_prefix} Insert failed. 'table' is set on InsertRequest (request-level), so "
+ "'upsert' must also be provided at the request level -- a record cannot set its own "
+ "'upsert' while 'table' is set at the request level."
+ )
+ REQUEST_LEVEL_UPSERT_NOT_ALLOWED_IN_INSERT = (
+ f"{error_prefix} Insert failed. 'table' is set per-record, so 'upsert' must also be "
+ "provided per-record -- InsertRequest's request-level 'upsert' cannot be used while "
+ "'table' is set on individual records."
+ )
+ EMPTY_KEY_IN_INSERT_DATA = f"{error_prefix} Insert failed. Each record's 'values' must not contain a null or empty key."
+ EMPTY_VALUE_IN_INSERT_DATA = f"{error_prefix} Insert failed. Each record's 'values' must not contain a null or empty value."
+
+ class Info(Enum):
+ VALIDATE_INSERT_REQUEST = f"{INFO}: [{error_prefix}] Validating insert request."
+ INSERT_TRIGGERED = f"{INFO}: [{error_prefix}] Insert method triggered."
+ INSERT_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Insert request resolved."
+ INSERT_SUCCESS = f"{INFO}: [{error_prefix}] Data inserted."
+
+ class ErrorLogs(Enum):
+ INSERT_RECORDS_REJECTED = f"{ERROR}: [{error_prefix}] Insert call resulted in failure."
diff --git a/flowvault/skyflow_flowvault/utils/_utils.py b/flowvault/skyflow_flowvault/utils/_utils.py
new file mode 100644
index 00000000..6bebb796
--- /dev/null
+++ b/flowvault/skyflow_flowvault/utils/_utils.py
@@ -0,0 +1,54 @@
+import platform
+import sys
+
+from common.errors import SkyflowError
+from common.utils import SkyflowMessages as CommonMessages
+from common.utils.constants import PROTOCOL, SdkMetricsKey, SdkPrefix
+from common.utils.enums import Env
+from ._version import SDK_VERSION
+from .enums import EnvUrls
+
+_CACHED_METRICS: dict = {}
+
+invalid_input_error_code = CommonMessages.ErrorCodes.INVALID_INPUT.value
+
+
+def get_vault_url(cluster_id, env, vault_id, logger=None):
+ """Mirrors common.utils.get_vault_url with v3's own EnvUrls (different subdomain)."""
+ if not cluster_id or not isinstance(cluster_id, str) or not cluster_id.strip():
+ raise SkyflowError(CommonMessages.Error.INVALID_CLUSTER_ID.value.format(vault_id), invalid_input_error_code)
+
+ if env not in Env:
+ raise SkyflowError(CommonMessages.Error.INVALID_ENV.value.format(vault_id), invalid_input_error_code)
+
+ base_url = EnvUrls[env.name].value
+
+ return f"{PROTOCOL}://{cluster_id}.{base_url}"
+
+
+def get_metrics():
+ if _CACHED_METRICS:
+ return _CACHED_METRICS
+
+ try:
+ sdk_client_device_model = platform.node()
+ except Exception:
+ sdk_client_device_model = ""
+
+ try:
+ sdk_client_os_details = sys.platform
+ except Exception:
+ sdk_client_os_details = ""
+
+ try:
+ sdk_runtime_details = sys.version
+ except Exception:
+ sdk_runtime_details = ""
+
+ _CACHED_METRICS.update({
+ SdkMetricsKey.SDK_NAME_VERSION: SdkPrefix.SKYFLOW_PYTHON + SDK_VERSION,
+ SdkMetricsKey.SDK_CLIENT_DEVICE_MODEL: sdk_client_device_model,
+ SdkMetricsKey.SDK_CLIENT_OS_DETAILS: sdk_client_os_details,
+ SdkMetricsKey.SDK_RUNTIME_DETAILS: SdkPrefix.PYTHON_RUNTIME + sdk_runtime_details,
+ })
+ return _CACHED_METRICS
diff --git a/flowvault/skyflow_flowvault/utils/_version.py b/flowvault/skyflow_flowvault/utils/_version.py
new file mode 100644
index 00000000..a7755720
--- /dev/null
+++ b/flowvault/skyflow_flowvault/utils/_version.py
@@ -0,0 +1 @@
+SDK_VERSION = '1.0.0'
diff --git a/flowvault/skyflow_flowvault/utils/enums/__init__.py b/flowvault/skyflow_flowvault/utils/enums/__init__.py
new file mode 100644
index 00000000..b9232bfc
--- /dev/null
+++ b/flowvault/skyflow_flowvault/utils/enums/__init__.py
@@ -0,0 +1,2 @@
+from ._upsert_type import UpsertType
+from ._env_urls import EnvUrls
diff --git a/flowvault/skyflow_flowvault/utils/enums/_env_urls.py b/flowvault/skyflow_flowvault/utils/enums/_env_urls.py
new file mode 100644
index 00000000..ccaa2919
--- /dev/null
+++ b/flowvault/skyflow_flowvault/utils/enums/_env_urls.py
@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+class EnvUrls(Enum):
+ """v3 vault hosts -- a different subdomain than v2 (skyvault vs. vault). All four confirmed."""
+ DEV = "skyvault.skyflowapis.dev"
+ PROD = "skyvault.skyflowapis.com"
+ SANDBOX = "skyvault.skyflowapis-preview.com"
+ STAGE = "skyvault.skyflowapis.tech"
diff --git a/flowvault/skyflow_flowvault/utils/enums/_upsert_type.py b/flowvault/skyflow_flowvault/utils/enums/_upsert_type.py
new file mode 100644
index 00000000..b2803796
--- /dev/null
+++ b/flowvault/skyflow_flowvault/utils/enums/_upsert_type.py
@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class UpsertType(Enum):
+ """Mirrors the wire enum FlowEnumUpdateType (V1Upsert.update_type)."""
+ REPLACE = "REPLACE"
+ UPDATE = "UPDATE"
diff --git a/flowvault/skyflow_flowvault/utils/validations/__init__.py b/flowvault/skyflow_flowvault/utils/validations/__init__.py
new file mode 100644
index 00000000..499bed61
--- /dev/null
+++ b/flowvault/skyflow_flowvault/utils/validations/__init__.py
@@ -0,0 +1 @@
+from ._validations import validate_vault_config, validate_update_vault_config, validate_insert_request
diff --git a/flowvault/skyflow_flowvault/utils/validations/_validations.py b/flowvault/skyflow_flowvault/utils/validations/_validations.py
new file mode 100644
index 00000000..a8ea74df
--- /dev/null
+++ b/flowvault/skyflow_flowvault/utils/validations/_validations.py
@@ -0,0 +1,79 @@
+from common.errors import SkyflowError
+from common.utils import SkyflowMessages as CommonMessages
+from common.utils.validations import (
+ validate_keys,
+ validate_credentials,
+ validate_vault_config,
+ validate_update_vault_config,
+)
+from skyflow_flowvault.utils import SkyflowMessages
+from skyflow_flowvault.utils.enums import UpsertType
+
+VALID_INSERT_RECORD_KEYS = ["values", "table", "upsert"]
+VALID_UPSERT_KEYS = ["update_type", "unique_columns"]
+
+invalid_input_error_code = CommonMessages.ErrorCodes.INVALID_INPUT.value
+
+# validate_vault_config/validate_update_vault_config/validate_credentials are re-exported
+# directly from common.utils.validations -- flowvault's own logic here was field-for-field
+# identical to v2's, confirmed, so both variants now share one implementation.
+
+
+def _validate_upsert(logger, upsert):
+ if upsert is None:
+ return
+ if not isinstance(upsert, dict):
+ raise SkyflowError(SkyflowMessages.Error.INVALID_UPSERT_TYPE_IN_INSERT.value, invalid_input_error_code)
+ validate_keys(logger, upsert, VALID_UPSERT_KEYS)
+ unique_columns = upsert.get("unique_columns")
+ if (not isinstance(unique_columns, list) or not unique_columns
+ or not all(isinstance(c, str) for c in unique_columns)):
+ raise SkyflowError(SkyflowMessages.Error.INVALID_UPSERT_UNIQUE_COLUMNS_IN_INSERT.value, invalid_input_error_code)
+ update_type = upsert.get("update_type")
+ if update_type is not None and not isinstance(update_type, UpsertType):
+ raise SkyflowError(SkyflowMessages.Error.INVALID_UPSERT_UPDATE_TYPE_IN_INSERT.value, invalid_input_error_code)
+
+
+MAX_INSERT_RECORDS = 10000 # matches Java's v3 Validations.validateInsertRequest (hardcoded, not configurable)
+
+
+def validate_insert_request(logger, request):
+ if not isinstance(request.values, list) or not all(isinstance(r, dict) for r in request.values):
+ raise SkyflowError(SkyflowMessages.Error.INVALID_RECORDS_TYPE_IN_INSERT.value, invalid_input_error_code)
+
+ if not request.values:
+ raise SkyflowError(SkyflowMessages.Error.EMPTY_RECORDS_IN_INSERT.value, invalid_input_error_code)
+
+ if len(request.values) > MAX_INSERT_RECORDS:
+ raise SkyflowError(SkyflowMessages.Error.TOO_MANY_RECORDS_IN_INSERT.value, invalid_input_error_code)
+
+ # request.table/record["table"] format and record["values"] emptiness/key/value validity are
+ # checked by the controller via the shared BaseVaultController._validate_table_name_if_present()
+ # / _validate_field_values() -- not here, to avoid duplicating that logic.
+
+ _validate_upsert(logger, request.upsert)
+
+ for record in request.values:
+ validate_keys(logger, record, VALID_INSERT_RECORD_KEYS)
+ _validate_upsert(logger, record.get("upsert"))
+
+ # table must be set in exactly one place -- request-level (every record) or per-record (no
+ # partial mix) -- and upsert must live at that same place (mirrors Java's v3 Validations).
+ table_at_request_level = request.table is not None
+
+ if table_at_request_level:
+ for record in request.values:
+ if record.get("table") is not None:
+ raise SkyflowError(SkyflowMessages.Error.TABLE_NAME_IN_BOTH_PLACES_IN_INSERT.value, invalid_input_error_code)
+ else:
+ for record in request.values:
+ if record.get("table") is None:
+ raise SkyflowError(SkyflowMessages.Error.TABLE_NAME_MISSING_IN_INSERT.value, invalid_input_error_code)
+
+ if table_at_request_level:
+ for record in request.values:
+ if record.get("upsert") is not None:
+ raise SkyflowError(SkyflowMessages.Error.RECORD_LEVEL_UPSERT_NOT_ALLOWED_IN_INSERT.value, invalid_input_error_code)
+ else:
+ if request.upsert is not None:
+ raise SkyflowError(SkyflowMessages.Error.REQUEST_LEVEL_UPSERT_NOT_ALLOWED_IN_INSERT.value, invalid_input_error_code)
diff --git a/tests/vault/client/__init__.py b/flowvault/skyflow_flowvault/vault/__init__.py
similarity index 100%
rename from tests/vault/client/__init__.py
rename to flowvault/skyflow_flowvault/vault/__init__.py
diff --git a/tests/vault/connection/__init__.py b/flowvault/skyflow_flowvault/vault/client/__init__.py
similarity index 100%
rename from tests/vault/connection/__init__.py
rename to flowvault/skyflow_flowvault/vault/client/__init__.py
diff --git a/flowvault/skyflow_flowvault/vault/client/client.py b/flowvault/skyflow_flowvault/vault/client/client.py
new file mode 100644
index 00000000..5dc4c47d
--- /dev/null
+++ b/flowvault/skyflow_flowvault/vault/client/client.py
@@ -0,0 +1,14 @@
+from common.vault.base_vault_client import BaseVaultClient
+from skyflow_flowvault.generated.rest.client import SkyflowAuth
+from skyflow_flowvault.utils import get_vault_url
+
+
+class VaultClient(BaseVaultClient):
+ def resolve_vault_url(self, cluster_id, env, vault_id, logger=None):
+ return get_vault_url(cluster_id, env, vault_id, logger=logger)
+
+ def initialize_api_client(self, vault_url, bearer_token):
+ self._api_client = SkyflowAuth(base_url=vault_url)
+
+ def get_insert_api(self):
+ return self._api_client.flowservice
diff --git a/flowvault/skyflow_flowvault/vault/controller/__init__.py b/flowvault/skyflow_flowvault/vault/controller/__init__.py
new file mode 100644
index 00000000..27660679
--- /dev/null
+++ b/flowvault/skyflow_flowvault/vault/controller/__init__.py
@@ -0,0 +1 @@
+from ._vault import VaultController
diff --git a/flowvault/skyflow_flowvault/vault/controller/_vault.py b/flowvault/skyflow_flowvault/vault/controller/_vault.py
new file mode 100644
index 00000000..a6abaa49
--- /dev/null
+++ b/flowvault/skyflow_flowvault/vault/controller/_vault.py
@@ -0,0 +1,158 @@
+import json
+
+from common.utils import SkyflowMessages as CommonMessages
+from common.utils.constants import SKY_META_DATA_HEADER
+from common.utils.logger import log_info, log_error_log
+from common.vault.base_vault_controller import BaseVaultController
+from skyflow_flowvault.generated.rest import V1InsertRecordData, V1Upsert
+from skyflow_flowvault.generated.rest.core import ApiError
+from skyflow_flowvault.utils import SkyflowMessages, get_metrics
+from skyflow_flowvault.utils.validations import validate_insert_request
+from skyflow_flowvault.vault.data import InsertRequest, InsertResponse
+
+REQUEST_ID_HEADER = "x-request-id"
+
+
+class VaultController(BaseVaultController):
+ _skyflow_messages = SkyflowMessages
+
+ def __init__(self, vault_client):
+ super().__init__(vault_client)
+
+ def insert(self, request: InsertRequest) -> InsertResponse:
+ log_info(SkyflowMessages.Info.VALIDATE_INSERT_REQUEST.value, self._vault_client.get_logger())
+ validate_insert_request(self._vault_client.get_logger(), request)
+ self._validate_table_name_if_present(request.table)
+ for record in request.values:
+ self._validate_table_name_if_present(record.get("table"))
+ self._validate_field_values(record.get("values"))
+ log_info(SkyflowMessages.Info.INSERT_REQUEST_RESOLVED.value, self._vault_client.get_logger())
+ self._vault_client.initialize_client_configuration()
+
+ insert_api = self._vault_client.get_insert_api()
+
+ needs_per_record_table = any(r.get("table") is not None for r in request.values)
+ needs_per_record_upsert = any(r.get("upsert") is not None for r in request.values)
+
+ wire_records = [
+ self.__build_wire_record(record, request, needs_per_record_table, needs_per_record_upsert)
+ for record in request.values
+ ]
+
+ try:
+ log_info(SkyflowMessages.Info.INSERT_TRIGGERED.value, self._vault_client.get_logger())
+ headers = self.__build_headers()
+ top_level_kwargs = self.__omit_none(
+ table_name=None if needs_per_record_table else request.table,
+ upsert=None if needs_per_record_upsert else self.__to_v1_upsert(request.upsert),
+ )
+ raw_response = insert_api.with_raw_response.insert(
+ vault_id=self._vault_client.get_vault_id(),
+ records=wire_records,
+ request_options={'additional_headers': headers},
+ **top_level_kwargs,
+ )
+ request_id = self.__extract_request_id(raw_response.headers)
+ inserted_fields, errors = self.__split_success_and_errors(raw_response.data.records or [], 0, request_id)
+ except Exception as e:
+ log_error_log(SkyflowMessages.ErrorLogs.INSERT_RECORDS_REJECTED.value, self._vault_client.get_logger())
+ inserted_fields, errors = [], self.__errors_from_exception(e, request.values, 0)
+
+ log_info(SkyflowMessages.Info.INSERT_SUCCESS.value, self._vault_client.get_logger())
+ return InsertResponse(inserted_fields=inserted_fields, errors=errors if errors else None)
+
+ def get(self, request):
+ raise NotImplementedError("VaultController.get is not implemented yet")
+
+ def update(self, request):
+ raise NotImplementedError("VaultController.update is not implemented yet")
+
+ def delete(self, request):
+ raise NotImplementedError("VaultController.delete is not implemented yet")
+
+ def query(self, request):
+ raise NotImplementedError("VaultController.query is not implemented yet")
+
+ def detokenize(self, request):
+ raise NotImplementedError("VaultController.detokenize is not implemented yet")
+
+ def __build_wire_record(self, record, request, needs_per_record_table, needs_per_record_upsert):
+ return V1InsertRecordData(data=record["values"], **self.__omit_none(
+ table_name=(record.get("table") or request.table) if needs_per_record_table else None,
+ upsert=self.__to_v1_upsert(record.get("upsert") or request.upsert) if needs_per_record_upsert else None,
+ ))
+
+ def __omit_none(self, **kwargs):
+ return {k: v for k, v in kwargs.items() if v is not None}
+
+ def __build_headers(self):
+ headers = {SKY_META_DATA_HEADER: json.dumps(get_metrics())}
+ token = self._vault_client.get_current_bearer_token()
+ if token:
+ headers['Authorization'] = f'Bearer {token}'
+ return headers
+
+ def __to_v1_upsert(self, upsert):
+ if upsert is None:
+ return None
+ update_type = upsert.get("update_type")
+ return V1Upsert(
+ update_type=update_type.value if update_type else None,
+ unique_columns=upsert.get("unique_columns"),
+ )
+
+ def __extract_request_id(self, headers):
+ return headers.get(REQUEST_ID_HEADER) if headers else None
+
+ def __split_success_and_errors(self, records, start_index, request_id):
+
+ successes, errors = [], []
+ for offset, record in enumerate(records):
+ request_index = start_index + offset
+ if record.error is not None:
+ errors.append({'request_index': request_index, 'error': record.error, 'code': record.http_code, 'request_id': request_id})
+ else:
+ success = {
+ 'request_index': request_index,
+ 'skyflow_id': record.skyflow_id,
+ }
+ success.update(self.__flatten_tokens(record.tokens))
+ successes.append(success)
+ return successes, errors
+
+ def __flatten_tokens(self, tokens):
+ if not tokens:
+ return {}
+ flat = {}
+ for column, entries in tokens.items():
+ if isinstance(entries, list):
+ token_values = [entry.get('token') for entry in entries if isinstance(entry, dict)]
+ flat[column] = token_values[0] if len(token_values) == 1 else token_values
+ else:
+ flat[column] = entries
+ return flat
+
+ def __errors_from_exception(self, e, records, start_index):
+ if isinstance(e, ApiError):
+ request_id = self.__extract_request_id(e.headers)
+ body = e.body if isinstance(e.body, dict) else None
+ if body and isinstance(body.get('records'), list) and body['records']:
+ return [
+ self.__error_dict_from_record_map(record, start_index + offset, request_id)
+ for offset, record in enumerate(body['records']) if isinstance(record, dict)
+ ]
+ if body and body.get('error') is not None:
+ err_field = body['error']
+ if isinstance(err_field, dict):
+ return [self.__error_dict_from_record_map(err_field, start_index + i, request_id) for i in range(len(records))]
+ return [{'request_index': start_index + i, 'error': str(err_field), 'code': e.status_code, 'request_id': request_id}
+ for i in range(len(records))]
+ return [{'request_index': start_index + i, 'error': str(e), 'code': e.status_code, 'request_id': request_id}
+ for i in range(len(records))]
+ message = str(e) if e else CommonMessages.Error.GENERIC_API_ERROR.value
+ return [{'request_index': start_index + i, 'error': message, 'code': None, 'request_id': None} for i in range(len(records))]
+
+ def __error_dict_from_record_map(self, record_map, request_index, request_id):
+ code = record_map.get('http_code', record_map.get('httpCode', record_map.get('statusCode')))
+ message = record_map.get('error', record_map.get('message', 'Unknown error'))
+ return {'request_index': request_index, 'error': message, 'code': code, 'request_id': request_id}
diff --git a/flowvault/skyflow_flowvault/vault/data/__init__.py b/flowvault/skyflow_flowvault/vault/data/__init__.py
new file mode 100644
index 00000000..62ae85cc
--- /dev/null
+++ b/flowvault/skyflow_flowvault/vault/data/__init__.py
@@ -0,0 +1,3 @@
+from ._insert_request import InsertRequest
+from ._insert_response import InsertResponse
+from ._upsert import Upsert
diff --git a/flowvault/skyflow_flowvault/vault/data/_insert_request.py b/flowvault/skyflow_flowvault/vault/data/_insert_request.py
new file mode 100644
index 00000000..a6da5c11
--- /dev/null
+++ b/flowvault/skyflow_flowvault/vault/data/_insert_request.py
@@ -0,0 +1,7 @@
+from common.vault.data import BaseInsertRequest
+from skyflow_flowvault.vault.data._upsert import Upsert
+
+
+class InsertRequest(BaseInsertRequest):
+ def __init__(self, values: list, table: str = None, upsert: Upsert = None):
+ super().__init__(table, values, upsert=upsert)
diff --git a/flowvault/skyflow_flowvault/vault/data/_insert_response.py b/flowvault/skyflow_flowvault/vault/data/_insert_response.py
new file mode 100644
index 00000000..ddb87134
--- /dev/null
+++ b/flowvault/skyflow_flowvault/vault/data/_insert_response.py
@@ -0,0 +1,6 @@
+from common.vault.data import BaseInsertResponse
+
+
+class InsertResponse(BaseInsertResponse):
+ """flowvault's own insert() response class -- currently identical to the shared base, kept as
+ its own subclass so flowvault-specific fields can be added later without touching PDB."""
diff --git a/flowvault/skyflow_flowvault/vault/data/_upsert.py b/flowvault/skyflow_flowvault/vault/data/_upsert.py
new file mode 100644
index 00000000..27d8fdab
--- /dev/null
+++ b/flowvault/skyflow_flowvault/vault/data/_upsert.py
@@ -0,0 +1,6 @@
+from typing import Optional, TypedDict
+from skyflow_flowvault.utils.enums import UpsertType
+
+class Upsert(TypedDict, total=False):
+ update_type: Optional[UpsertType]
+ unique_columns: list
diff --git a/tests/vault/controller/__init__.py b/flowvault/tests/__init__.py
similarity index 100%
rename from tests/vault/controller/__init__.py
rename to flowvault/tests/__init__.py
diff --git a/tests/vault/data/__init__.py b/flowvault/tests/client/__init__.py
similarity index 100%
rename from tests/vault/data/__init__.py
rename to flowvault/tests/client/__init__.py
diff --git a/flowvault/tests/client/test_skyflow.py b/flowvault/tests/client/test_skyflow.py
new file mode 100644
index 00000000..c7d908f2
--- /dev/null
+++ b/flowvault/tests/client/test_skyflow.py
@@ -0,0 +1,130 @@
+import unittest
+from unittest.mock import patch
+
+from common.errors import SkyflowError
+from common.utils import LogLevel, Env
+from skyflow_flowvault.client import Skyflow
+
+VALID_VAULT_CONFIG = {
+ "vault_id": "VAULT_ID",
+ "cluster_id": "CLUSTER_ID",
+ "env": Env.DEV,
+ "credentials": {"path": "/path/to/valid_credentials.json"},
+}
+
+INVALID_VAULT_CONFIG = {
+ "cluster_id": "CLUSTER_ID", # missing vault_id
+ "env": Env.DEV,
+ "credentials": {"path": "/path/to/valid_credentials.json"},
+}
+
+VALID_CREDENTIALS = {"path": "/path/to/valid_credentials.json"}
+
+
+class TestSkyflowVaultConfig(unittest.TestCase):
+ """v2 parity: flowvault gained remove_vault_config/update_vault_config/
+ update_skyflow_credentials/update_log_level/get_log_level via the shared common base --
+ these confirm they actually work here too, not just on v2."""
+
+ def setUp(self):
+ self.builder = Skyflow.builder()
+
+ def test_add_vault_config_success(self):
+ builder = self.builder.add_vault_config(VALID_VAULT_CONFIG)
+ self.assertEqual(builder, self.builder)
+
+ def test_add_vault_config_invalid_raises(self):
+ with self.assertRaises(SkyflowError):
+ self.builder.add_vault_config(INVALID_VAULT_CONFIG)
+
+ def test_build_and_get_vault_config(self):
+ client = self.builder.add_vault_config(VALID_VAULT_CONFIG).build()
+ config = client.get_vault_config("VAULT_ID")
+ self.assertEqual(config.get("vault_id"), "VAULT_ID")
+
+ @patch("skyflow_flowvault.vault.client.client.VaultClient.update_config")
+ def test_update_vault_config(self, mock_update_config):
+ client = self.builder.add_vault_config(VALID_VAULT_CONFIG).build()
+ updated = dict(VALID_VAULT_CONFIG)
+ updated["cluster_id"] = "NEW_CLUSTER"
+ client.update_vault_config(updated)
+ mock_update_config.assert_called_once()
+
+ def test_update_vault_config_with_invalid_vault_id_raises(self):
+ client = self.builder.add_vault_config(VALID_VAULT_CONFIG).build()
+ invalid = dict(VALID_VAULT_CONFIG)
+ invalid["vault_id"] = "does_not_exist"
+ with self.assertRaises(SkyflowError):
+ client.update_vault_config(invalid)
+
+ def test_remove_vault_config(self):
+ client = self.builder.add_vault_config(VALID_VAULT_CONFIG).build()
+ client.remove_vault_config("VAULT_ID")
+ with self.assertRaises(SkyflowError):
+ client.get_vault_config("VAULT_ID")
+
+ def test_add_and_update_skyflow_credentials(self):
+ client = self.builder.add_vault_config(VALID_VAULT_CONFIG).build()
+ client.add_skyflow_credentials(VALID_CREDENTIALS)
+ new_credentials = {"path": "/path/to/other_credentials.json"}
+ client.update_skyflow_credentials(new_credentials)
+ # no assertion error means both went through the same underlying builder path
+
+ def test_set_and_get_log_level(self):
+ client = self.builder.add_vault_config(VALID_VAULT_CONFIG).build()
+ client.set_log_level(LogLevel.INFO)
+ self.assertEqual(client.get_log_level(), LogLevel.INFO)
+
+ def test_update_log_level_delegates_to_set_log_level(self):
+ client = self.builder.add_vault_config(VALID_VAULT_CONFIG).build()
+ client.update_log_level(LogLevel.INFO)
+ self.assertEqual(client.get_log_level(), LogLevel.INFO)
+
+ @patch("common.client.base_skyflow.log_warn")
+ def test_update_log_level_emits_deprecation_warning(self, mock_warn):
+ client = self.builder.add_vault_config(VALID_VAULT_CONFIG).build()
+ client.update_log_level(LogLevel.INFO)
+ mock_warn.assert_called_once()
+ self.assertIn("set_log_level", mock_warn.call_args[0][0])
+
+ def test_vault_returns_vault_controller(self):
+ client = self.builder.add_vault_config(VALID_VAULT_CONFIG).build()
+ vault = client.vault("VAULT_ID")
+ self.assertTrue(hasattr(vault, "insert"))
+
+
+class TestConnectionAndDetectNotSupported(unittest.TestCase):
+ """flowvault has no Connection/Detect concept this round -- ConnectionMixin/DetectMixin are
+ only added to a variant's produced Skyflow class when connection_cls/detect_cls are supplied,
+ so these methods are genuinely absent here (AttributeError), unlike v2 where they work."""
+
+ def setUp(self):
+ self.client = Skyflow.builder().add_vault_config(VALID_VAULT_CONFIG).build()
+
+ def test_connection_raises(self):
+ with self.assertRaises(AttributeError):
+ self.client.connection()
+
+ def test_detect_raises(self):
+ with self.assertRaises(AttributeError):
+ self.client.detect()
+
+ def test_add_connection_config_raises(self):
+ with self.assertRaises(AttributeError):
+ self.client.add_connection_config({})
+
+ def test_remove_connection_config_raises(self):
+ with self.assertRaises(AttributeError):
+ self.client.remove_connection_config("x")
+
+ def test_update_connection_config_raises(self):
+ with self.assertRaises(AttributeError):
+ self.client.update_connection_config({})
+
+ def test_get_connection_config_raises(self):
+ with self.assertRaises(AttributeError):
+ self.client.get_connection_config("x")
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/vault/detect/__init__.py b/flowvault/tests/utils/__init__.py
similarity index 100%
rename from tests/vault/detect/__init__.py
rename to flowvault/tests/utils/__init__.py
diff --git a/tests/vault/tokens/__init__.py b/flowvault/tests/utils/validations/__init__.py
similarity index 100%
rename from tests/vault/tokens/__init__.py
rename to flowvault/tests/utils/validations/__init__.py
diff --git a/flowvault/tests/utils/validations/test__validations.py b/flowvault/tests/utils/validations/test__validations.py
new file mode 100644
index 00000000..10b36b6b
--- /dev/null
+++ b/flowvault/tests/utils/validations/test__validations.py
@@ -0,0 +1,182 @@
+import unittest
+
+from common.errors import SkyflowError
+from common.utils.enums import Env
+from skyflow_flowvault.utils.enums import UpsertType
+from skyflow_flowvault.utils.validations import validate_insert_request, validate_vault_config
+from skyflow_flowvault.vault.data import InsertRequest
+
+
+class TestValidateInsertRequest(unittest.TestCase):
+ def test_valid_minimal_request(self):
+ request = InsertRequest(values=[dict(values={"a": 1})], table="t1")
+ validate_insert_request(None, request) # should not raise
+
+ def test_valid_rich_request_with_per_record_overrides(self):
+ """Valid per-record use: no request-level table/upsert at all (the vault rejects
+ setting it in both places -- see test_table_in_both_places_raises below). Java parity
+ requires EVERY record to set its own table when there's no request-level table (a
+ partial mix is invalid -- see test_table_missing_from_one_record_raises), so both
+ records set their own here."""
+ request = InsertRequest(
+ values=[
+ dict(values={"a": 1}, table="t2"),
+ dict(values={"a": 2}, table="t2", upsert={"update_type": UpsertType.REPLACE, "unique_columns": ["a"]}),
+ ],
+ )
+ validate_insert_request(None, request) # should not raise
+
+ def test_table_in_both_places_raises(self):
+ """Confirmed directly against a real vault: 'Table name should be present outside the
+ records or inside each record. Should be present at one place.'"""
+ request = InsertRequest(
+ values=[dict(values={"a": 1}, table="t2")],
+ table="t1",
+ )
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+ def test_table_in_both_places_raises_even_if_only_one_record_sets_it(self):
+ request = InsertRequest(
+ values=[dict(values={"a": 1}, table="t2"), dict(values={"a": 2})],
+ table="t1",
+ )
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+ def test_record_level_upsert_forbidden_when_table_is_at_request_level(self):
+ """Java parity: 'upsert' must live at the SAME place as 'table'. Here table is at the
+ request level, so a record-level upsert is rejected even though this record's own table
+ placement (none) is fine."""
+ request = InsertRequest(
+ values=[dict(values={"a": 1}, upsert={"unique_columns": ["b"]})],
+ table="t1",
+ upsert={"unique_columns": ["a"]},
+ )
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+ def test_request_level_upsert_forbidden_when_table_is_per_record(self):
+ request = InsertRequest(
+ values=[dict(values={"a": 1}, table="t1")],
+ upsert={"unique_columns": ["a"]},
+ )
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+ def test_too_many_records_raises(self):
+ request = InsertRequest(values=[dict(values={"a": 1}) for _ in range(10001)], table="t1")
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+ def test_exactly_max_records_is_valid(self):
+ request = InsertRequest(values=[dict(values={"a": 1}) for _ in range(10000)], table="t1")
+ validate_insert_request(None, request) # should not raise
+
+ def test_table_missing_from_one_record_raises(self):
+ """Java parity: when there's no request-level table, EVERY record must set its own --
+ a partial mix (some records with a table, some without) is invalid."""
+ request = InsertRequest(values=[dict(values={"a": 1}, table="t1"), dict(values={"a": 2})])
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+ # Empty/null key or value in a record's 'values', and record 'values' being a non-empty
+ # dict, are now validated by the controller via the shared
+ # BaseVaultController._validate_field_values() -- see test__vault.py's
+ # test_insert_raises_on_empty_key/_on_empty_value/_on_non_dict_values/_on_empty_values, and
+ # common/tests/vault/test_base_vault_controller.py for the shared helper's own unit tests.
+
+ def test_falsy_non_string_values_are_valid(self):
+ """0, False, [], {} are all legitimate values -- only None/empty-string should raise
+ (mirrors Java's value.toString().trim().isEmpty(), which is non-empty for all of these)."""
+ request = InsertRequest(values=[dict(values={"a": 0, "b": False, "c": [], "d": {}})], table="t1")
+ validate_insert_request(None, request) # should not raise
+
+ def test_request_level_table_alone_is_valid(self):
+ request = InsertRequest(values=[dict(values={"a": 1}), dict(values={"a": 2})], table="t1")
+ validate_insert_request(None, request) # should not raise
+
+ def test_per_record_table_alone_is_valid(self):
+ request = InsertRequest(values=[dict(values={"a": 1}, table="t1"), dict(values={"a": 2}, table="t2")])
+ validate_insert_request(None, request) # should not raise
+
+ def test_records_must_be_a_list(self):
+ request = InsertRequest(values="not-a-list", table="t1")
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+ def test_records_must_be_dicts(self):
+ request = InsertRequest(values=["not-a-dict"], table="t1")
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+ def test_record_with_unknown_key_raises(self):
+ request = InsertRequest(values=[{"a": 1}], table="t1")
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+ def test_records_must_not_be_empty(self):
+ request = InsertRequest(values=[], table="t1")
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+ # Table name format (non-empty string if provided) is now validated by the controller via
+ # the shared BaseVaultController._validate_table_name_if_present() -- see
+ # test__vault.py's test_insert_raises_on_invalid_table_name and
+ # common/tests/vault/test_base_vault_controller.py for the shared helper's own unit tests.
+
+ def test_table_is_optional_when_every_record_has_its_own(self):
+ request = InsertRequest(values=[dict(values={"a": 1}, table="t2")])
+ validate_insert_request(None, request) # should not raise
+
+ def test_upsert_must_be_a_dict(self):
+ request = InsertRequest(values=[dict(values={"a": 1})], table="t1", upsert="not-an-upsert")
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+ def test_upsert_unique_columns_must_be_non_empty_list_of_strings(self):
+ request = InsertRequest(values=[dict(values={"a": 1})], table="t1", upsert={"unique_columns": []})
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+ def test_upsert_update_type_must_be_upsert_type_enum(self):
+ request = InsertRequest(
+ values=[dict(values={"a": 1})], table="t1",
+ upsert={"update_type": "REPLACE", "unique_columns": ["a"]}, # plain string, not the enum
+ )
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+ def test_per_record_upsert_is_also_validated(self):
+ request = InsertRequest(
+ values=[dict(values={"a": 1}, upsert={"unique_columns": []})],
+ table="t1",
+ )
+ with self.assertRaises(SkyflowError):
+ validate_insert_request(None, request)
+
+
+class TestValidateVaultConfig(unittest.TestCase):
+ def test_valid_config(self):
+ config = {
+ "vault_id": "vault123",
+ "cluster_id": "cluster1",
+ "env": Env.PROD,
+ # api_key (not a JWT-format "token") avoids the expiry check so this only
+ # exercises validate_vault_config's own structural validation.
+ "credentials": {"api_key": "sky-abcde-" + "f" * 32},
+ }
+ self.assertTrue(validate_vault_config(None, config))
+
+ def test_missing_vault_id_raises(self):
+ with self.assertRaises(SkyflowError):
+ validate_vault_config(None, {"cluster_id": "cluster1"})
+
+ def test_unknown_key_raises(self):
+ config = {"vault_id": "v", "cluster_id": "c", "unexpected_key": True}
+ with self.assertRaises(SkyflowError):
+ validate_vault_config(None, config)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/flowvault/tests/vault/__init__.py b/flowvault/tests/vault/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/flowvault/tests/vault/client/__init__.py b/flowvault/tests/vault/client/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/flowvault/tests/vault/client/test__client.py b/flowvault/tests/vault/client/test__client.py
new file mode 100644
index 00000000..2accba2b
--- /dev/null
+++ b/flowvault/tests/vault/client/test__client.py
@@ -0,0 +1,55 @@
+import unittest
+from unittest.mock import patch, MagicMock
+
+from common.utils.enums import Env
+from common.vault.base_vault_client import BaseVaultClient
+from skyflow_flowvault.vault.client.client import VaultClient
+
+
+class TestVaultClient(unittest.TestCase):
+ def setUp(self):
+ self.vault_client = VaultClient({"vault_id": "test_vault"})
+
+ def test_is_a_base_vault_client(self):
+ self.assertIsInstance(self.vault_client, BaseVaultClient)
+
+ # ------------------------------------------------------------------ #
+ # resolve_vault_url — v3's own domain (skyvault.skyflowapis.*), confirmed
+ # to differ from v2's (vault.skyflowapis.*) for the same cluster_id/env
+ # -- reusing v2's derivation here 404'd. All four envs confirmed.
+ # ------------------------------------------------------------------ #
+
+ def test_resolve_vault_url_uses_v3_skyvault_domain_dev(self):
+ url = self.vault_client.resolve_vault_url("qhdmceurtnlz", Env.DEV, "myvault")
+ self.assertEqual(url, "https://qhdmceurtnlz.skyvault.skyflowapis.dev")
+
+ def test_resolve_vault_url_uses_v3_skyvault_domain_prod(self):
+ url = self.vault_client.resolve_vault_url("qhdmceurtnlz", Env.PROD, "myvault")
+ self.assertEqual(url, "https://qhdmceurtnlz.skyvault.skyflowapis.com")
+
+ def test_resolve_vault_url_uses_v3_skyvault_domain_sandbox(self):
+ url = self.vault_client.resolve_vault_url("qhdmceurtnlz", Env.SANDBOX, "myvault")
+ self.assertEqual(url, "https://qhdmceurtnlz.skyvault.skyflowapis-preview.com")
+
+ def test_resolve_vault_url_uses_v3_skyvault_domain_stage(self):
+ url = self.vault_client.resolve_vault_url("qhdmceurtnlz", Env.STAGE, "myvault")
+ self.assertEqual(url, "https://qhdmceurtnlz.skyvault.skyflowapis.tech")
+
+ @patch("skyflow_flowvault.vault.client.client.SkyflowAuth")
+ def test_initialize_api_client_does_not_pass_token(self, mock_skyflow_auth):
+ """v3's generated client has no `token` param at all -- unlike v2, nothing should be
+ baked in at construction time; auth is injected per-call instead (see Vault._build_headers)."""
+ self.vault_client.initialize_api_client("https://test-vault-url.com", "some_bearer_token")
+
+ _, kwargs = mock_skyflow_auth.call_args
+ self.assertEqual(kwargs.get("base_url"), "https://test-vault-url.com")
+ self.assertNotIn("token", kwargs)
+
+ def test_get_insert_api_returns_flowservice(self):
+ self.vault_client._api_client = MagicMock()
+ result = self.vault_client.get_insert_api()
+ self.assertEqual(result, self.vault_client._api_client.flowservice)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/flowvault/tests/vault/controller/__init__.py b/flowvault/tests/vault/controller/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/flowvault/tests/vault/controller/test__vault.py b/flowvault/tests/vault/controller/test__vault.py
new file mode 100644
index 00000000..449ede9d
--- /dev/null
+++ b/flowvault/tests/vault/controller/test__vault.py
@@ -0,0 +1,375 @@
+import unittest
+from unittest.mock import MagicMock, Mock, patch
+
+from common.errors import SkyflowError
+from skyflow_flowvault.generated.rest.core import ApiError
+from skyflow_flowvault.vault.controller import VaultController
+from skyflow_flowvault.vault.data import InsertRequest
+from skyflow_flowvault.utils.enums import UpsertType
+
+
+class FakeRecordResponseObject:
+ def __init__(self, skyflow_id=None, tokens=None, data=None, error=None, http_code=None, table_name=None):
+ self.skyflow_id = skyflow_id
+ self.tokens = tokens
+ self.data = data
+ self.error = error
+ self.http_code = http_code
+ self.table_name = table_name
+
+
+class FakeV1InsertResponse:
+ def __init__(self, records):
+ self.records = records
+
+
+class FakeRawResponse:
+ """Stands in for the HttpResponse wrapper returned by with_raw_response.insert(...) --
+ exposes .data (the parsed V1InsertResponse) and .headers, mirroring the real generated
+ client's RawFlowserviceClient."""
+
+ def __init__(self, records, headers=None):
+ self.data = FakeV1InsertResponse(records)
+ self.headers = headers or {}
+
+
+class TestVault(unittest.TestCase):
+ def setUp(self):
+ self.vault_client = Mock()
+ self.vault_client.get_vault_id.return_value = "vault123"
+ self.vault_client.get_logger.return_value = Mock()
+ self.vault_client.get_current_bearer_token.return_value = None
+ self.insert_api = MagicMock()
+ self.vault_client.get_insert_api.return_value = self.insert_api
+ self.vault = VaultController(self.vault_client)
+
+ # ------------------------------------------------------------------ #
+ # validation / initialization sequencing
+ # ------------------------------------------------------------------ #
+
+ @patch("skyflow_flowvault.vault.controller._vault.validate_insert_request")
+ def test_insert_validates_before_initializing_client(self, mock_validate):
+ self.insert_api.with_raw_response.insert.return_value = FakeRawResponse([])
+ request = InsertRequest(values=[dict(values={"a": 1})], table="t1")
+
+ self.vault.insert(request)
+
+ mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
+ self.vault_client.initialize_client_configuration.assert_called_once()
+
+ def test_insert_raises_for_invalid_request(self):
+ with self.assertRaises(SkyflowError):
+ self.vault.insert(InsertRequest(values=[], table="t1"))
+ self.vault_client.initialize_client_configuration.assert_not_called()
+
+ # ------------------------------------------------------------------ #
+ # shared BaseVaultController validation helpers, exercised end-to-end via insert()
+ # (unit-tested in isolation in common/tests/vault/test_base_vault_controller.py)
+ # ------------------------------------------------------------------ #
+
+ def test_insert_raises_on_empty_key(self):
+ with self.assertRaises(SkyflowError):
+ self.vault.insert(InsertRequest(values=[dict(values={"": "value"})], table="t1"))
+ self.insert_api.with_raw_response.insert.assert_not_called()
+
+ def test_insert_allows_empty_value(self):
+ self.insert_api.with_raw_response.insert.return_value = FakeRawResponse([])
+ self.vault.insert(InsertRequest(values=[dict(values={"a": ""})], table="t1"))
+ self.insert_api.with_raw_response.insert.assert_called_once()
+
+ def test_insert_raises_on_non_dict_values(self):
+ with self.assertRaises(SkyflowError):
+ self.vault.insert(InsertRequest(values=[dict(values=["not", "a", "dict"])], table="t1"))
+
+ def test_insert_raises_on_empty_values_dict(self):
+ with self.assertRaises(SkyflowError):
+ self.vault.insert(InsertRequest(values=[dict(values={})], table="t1"))
+
+ def test_insert_raises_on_invalid_request_level_table_name(self):
+ with self.assertRaises(SkyflowError):
+ self.vault.insert(InsertRequest(values=[dict(values={"a": 1})], table=" "))
+
+ def test_insert_raises_on_invalid_per_record_table_name(self):
+ with self.assertRaises(SkyflowError):
+ self.vault.insert(InsertRequest(values=[dict(values={"a": 1}, table=" ")]))
+
+ # ------------------------------------------------------------------ #
+ # request -> wire field mapping
+ # ------------------------------------------------------------------ #
+
+ def test_maps_request_level_table_and_upsert(self):
+ """When no record sets its own table/upsert, both go ONLY at the request level -- the
+ vault rejects sending table_name/upsert in both places (see the validation tests), so
+ the wire records must NOT also carry a resolved copy."""
+ self.insert_api.with_raw_response.insert.return_value = FakeRawResponse([])
+ request = InsertRequest(
+ values=[dict(values={"a": 1})],
+ table="t1",
+ upsert={"update_type": UpsertType.REPLACE, "unique_columns": ["a"]},
+ )
+
+ self.vault.insert(request)
+
+ _, kwargs = self.insert_api.with_raw_response.insert.call_args
+ self.assertEqual(kwargs["vault_id"], "vault123")
+ self.assertEqual(kwargs["table_name"], "t1")
+ self.assertEqual(len(kwargs["records"]), 1)
+ self.assertEqual(kwargs["records"][0].data, {"a": 1})
+ self.assertIsNone(kwargs["records"][0].table_name) # NOT resolved onto the record
+ self.assertEqual(kwargs["upsert"].update_type, "REPLACE")
+ self.assertEqual(kwargs["upsert"].unique_columns, ["a"])
+
+ def test_setting_table_at_both_request_and_record_level_raises(self):
+ """The vault rejects table_name in both places at once -- confirmed directly against a
+ real vault. validate_insert_request (tested separately) is what actually raises this;
+ this test just confirms insert() surfaces it rather than silently choosing one."""
+ request = InsertRequest(
+ values=[dict(values={"a": 1}, table="t2")],
+ table="t1",
+ )
+
+ with self.assertRaises(SkyflowError):
+ self.vault.insert(request)
+ self.insert_api.with_raw_response.insert.assert_not_called()
+
+ def test_per_record_table_and_upsert_used_when_request_level_unset(self):
+ """Legitimate per-record use: no request-level table/upsert at all -- Java parity
+ requires EVERY record to set its own table in this mode (see validation tests), so both
+ records do; only the second also sets its own upsert."""
+ self.insert_api.with_raw_response.insert.return_value = FakeRawResponse([])
+ request = InsertRequest(values=[
+ dict(values={"a": 1}, table="t2", upsert={"unique_columns": ["b"]}),
+ dict(values={"a": 2}, table="t2"),
+ ])
+
+ self.vault.insert(request)
+
+ _, kwargs = self.insert_api.with_raw_response.insert.call_args
+ self.assertNotIn("table_name", kwargs)
+ self.assertNotIn("upsert", kwargs)
+ self.assertEqual(kwargs["records"][0].table_name, "t2")
+ self.assertEqual(kwargs["records"][0].upsert.unique_columns, ["b"])
+ self.assertEqual(kwargs["records"][1].table_name, "t2")
+ self.assertIsNone(kwargs["records"][1].upsert)
+
+ def test_no_request_level_table_is_omitted_not_sent_as_none(self):
+ self.insert_api.with_raw_response.insert.return_value = FakeRawResponse([])
+ request = InsertRequest(values=[dict(values={"a": 1}, table="t2")]) # no request-level table
+
+ self.vault.insert(request)
+
+ _, kwargs = self.insert_api.with_raw_response.insert.call_args
+ self.assertNotIn("table_name", kwargs)
+
+ def test_wire_shape_matches_confirmed_working_request(self):
+ """Regression pin for a real bug: a request with only per-record table/upsert (no
+ request-level table/upsert at all) previously sent explicit `"tableName": null` /
+ `"upsert": null` at the top level, which diverged from a hand-verified working request
+ against a real vault (confirmed to have neither key present when unset)."""
+ self.insert_api.with_raw_response.insert.return_value = FakeRawResponse([])
+ request = InsertRequest(values=[
+ dict(
+ values={"name": "saileshwar", "email": "nanana@gmail.com"},
+ table="table1",
+ upsert={"update_type": UpsertType.UPDATE, "unique_columns": ["email"]},
+ ),
+ ])
+
+ self.vault.insert(request)
+
+ _, kwargs = self.insert_api.with_raw_response.insert.call_args
+ self.assertNotIn("table_name", kwargs)
+ self.assertNotIn("upsert", kwargs)
+ self.assertEqual(kwargs["records"][0].table_name, "table1")
+ self.assertEqual(kwargs["records"][0].upsert.update_type, "UPDATE")
+ self.assertEqual(kwargs["records"][0].upsert.unique_columns, ["email"])
+
+ def test_no_upsert_is_omitted_not_sent_as_none(self):
+ """upsert must be OMITTED from the wire call entirely when unset, not passed as None --
+ a real vault confirmed a working request never includes a null upsert/tableName key."""
+ self.insert_api.with_raw_response.insert.return_value = FakeRawResponse([])
+ request = InsertRequest(values=[dict(values={"a": 1})], table="t1")
+
+ self.vault.insert(request)
+
+ _, kwargs = self.insert_api.with_raw_response.insert.call_args
+ self.assertNotIn("upsert", kwargs)
+ self.assertIsNone(kwargs["records"][0].upsert)
+
+ # ------------------------------------------------------------------ #
+ # response shape -- mirrors PDB's InsertResponse (inserted_fields/errors)
+ # ------------------------------------------------------------------ #
+
+ def test_successful_records_go_to_inserted_fields(self):
+ self.insert_api.with_raw_response.insert.return_value = FakeRawResponse([
+ FakeRecordResponseObject(
+ skyflow_id="id1",
+ tokens={"name": [{"token": "tok1", "tokenGroupName": "deterministic_string"}]},
+ data={"name": "john doe"},
+ table_name="table1",
+ ),
+ ], headers={"x-request-id": "req-1"})
+ response = self.vault.insert(InsertRequest(values=[dict(values={"name": "john doe"})], table="table1"))
+
+ self.assertEqual(len(response.inserted_fields), 1)
+ inserted = response.inserted_fields[0]
+ self.assertEqual(inserted["request_index"], 0)
+ self.assertEqual(inserted["skyflow_id"], "id1")
+ self.assertEqual(inserted["name"], "tok1")
+ self.assertNotIn("data", inserted)
+ self.assertNotIn("table", inserted)
+ self.assertNotIn("tokens", inserted)
+ self.assertIsNone(response.errors)
+
+ def test_multiple_token_groups_for_one_field_flatten_to_a_list(self):
+ self.insert_api.with_raw_response.insert.return_value = FakeRawResponse([
+ FakeRecordResponseObject(
+ skyflow_id="id1",
+ tokens={"email": [
+ {"token": "tok-det", "tokenGroupName": "deterministic_string"},
+ {"token": "tok-nondet", "tokenGroupName": "nondeterministic_string"},
+ ]},
+ ),
+ ])
+ response = self.vault.insert(InsertRequest(values=[dict(values={"email": "a@b.com"})], table="t1"))
+
+ self.assertEqual(response.inserted_fields[0]["email"], ["tok-det", "tok-nondet"])
+
+ def test_mixed_success_and_error_records_are_split(self):
+ self.insert_api.with_raw_response.insert.return_value = FakeRawResponse([
+ FakeRecordResponseObject(skyflow_id="id1", tokens=None),
+ FakeRecordResponseObject(error="bad row", http_code=400, table_name="t1"),
+ ], headers={"x-request-id": "req-2"})
+ response = self.vault.insert(InsertRequest(
+ values=[dict(values={"a": 1}), dict(values={"a": 2})], table="t1",
+ ))
+
+ self.assertEqual(len(response.inserted_fields), 1)
+ self.assertEqual(response.inserted_fields[0]["request_index"], 0)
+ self.assertEqual(response.inserted_fields[0]["skyflow_id"], "id1")
+ self.assertEqual(len(response.errors), 1)
+ self.assertEqual(response.errors[0]["request_index"], 1)
+ self.assertEqual(response.errors[0]["error"], "bad row")
+ self.assertEqual(response.errors[0]["code"], 400)
+ self.assertEqual(response.errors[0]["request_id"], "req-2")
+
+ def test_error_record_identified_by_error_field_alone(self):
+ """Mirrors Java's Utils.formatResponse exactly: a record is an error purely by .error
+ being present -- http_code is read onto the error dict's 'code' key but is not itself
+ part of the success/error decision."""
+ self.insert_api.with_raw_response.insert.return_value = FakeRawResponse([
+ FakeRecordResponseObject(skyflow_id="id1", http_code=200),
+ ])
+ response = self.vault.insert(InsertRequest(values=[dict(values={"a": 1})], table="t1"))
+
+ self.assertEqual(len(response.inserted_fields), 1)
+ self.assertIsNone(response.errors)
+
+ # ------------------------------------------------------------------ #
+ # no batching -- every insert is exactly one API call
+ # ------------------------------------------------------------------ #
+
+ def test_all_records_sent_in_a_single_api_call_regardless_of_count(self):
+ self.insert_api.with_raw_response.insert.side_effect = lambda **kwargs: FakeRawResponse(
+ [FakeRecordResponseObject(skyflow_id=f"id-{i}") for i in range(len(kwargs["records"]))]
+ )
+ records = [dict(values={"a": i}) for i in range(4)]
+
+ response = self.vault.insert(InsertRequest(values=records, table="t1"))
+
+ self.insert_api.with_raw_response.insert.assert_called_once()
+ call_size = len(self.insert_api.with_raw_response.insert.call_args.kwargs["records"])
+ self.assertEqual(call_size, 4)
+ self.assertEqual(len(response.inserted_fields), 4)
+
+ def test_request_index_matches_position_in_the_original_records_list(self):
+ self.insert_api.with_raw_response.insert.side_effect = lambda **kwargs: FakeRawResponse(
+ [FakeRecordResponseObject(skyflow_id=f"id-{i}") for i in range(len(kwargs["records"]))]
+ )
+ records = [dict(values={"a": i}) for i in range(4)]
+
+ response = self.vault.insert(InsertRequest(values=records, table="t1"))
+
+ self.assertEqual(sorted(s["request_index"] for s in response.inserted_fields), [0, 1, 2, 3])
+
+ # ------------------------------------------------------------------ #
+ # transport failure
+ # ------------------------------------------------------------------ #
+
+ def test_transport_exception_marks_every_record_as_an_error(self):
+ """Without batching, one API call carries every record -- a transport-level exception
+ on that single call means every record in the request fails, not just some."""
+ self.insert_api.with_raw_response.insert.side_effect = Exception("network blip")
+ records = [dict(values={"a": 1}), dict(values={"a": 2})]
+
+ response = self.vault.insert(InsertRequest(values=records, table="t1"))
+
+ self.insert_api.with_raw_response.insert.assert_called_once()
+ self.assertEqual(len(response.inserted_fields), 0)
+ self.assertEqual(len(response.errors), 2)
+ self.assertTrue(all("network blip" in e["error"] for e in response.errors))
+ self.assertEqual([e["request_index"] for e in response.errors], [0, 1])
+
+ def test_api_error_with_structured_per_record_body_splits_into_one_error_per_row(self):
+ """Mirrors Java's Utils.handleBatchException: a structured error body (a 'records' list)
+ is split into individual error dicts instead of repeating one flat message for the
+ whole batch -- shaped after a real vault's actual 400 response for a partial-batch
+ failure (e.g. a NOT NULL column violation on one row)."""
+ api_error = ApiError(
+ status_code=400,
+ headers={"x-request-id": "req-3"},
+ body={"records": [
+ {"error": "Column passport has the notNull attribute, and input contains a null value.",
+ "httpCode": 400},
+ ]},
+ )
+ self.insert_api.with_raw_response.insert.side_effect = api_error
+
+ response = self.vault.insert(InsertRequest(values=[dict(values={"name": "a"})], table="t1"))
+
+ self.assertEqual(len(response.errors), 1)
+ self.assertIn("notNull", response.errors[0]["error"])
+ self.assertEqual(response.errors[0]["code"], 400)
+ self.assertEqual(response.errors[0]["request_id"], "req-3")
+ self.assertEqual(response.errors[0]["request_index"], 0)
+
+ def test_api_error_with_flat_body_falls_back_to_one_error_per_record(self):
+ api_error = ApiError(status_code=500, headers={}, body={"error": "internal error"})
+ self.insert_api.with_raw_response.insert.side_effect = api_error
+
+ response = self.vault.insert(InsertRequest(
+ values=[dict(values={"a": 1}), dict(values={"a": 2})], table="t1",
+ ))
+
+ self.assertEqual(len(response.errors), 2)
+ self.assertTrue(all(e["error"] == "internal error" for e in response.errors))
+ self.assertTrue(all(e["code"] == 500 for e in response.errors))
+ self.assertEqual([e["request_index"] for e in response.errors], [0, 1])
+
+ # ------------------------------------------------------------------ #
+ # per-call Authorization header injection
+ # ------------------------------------------------------------------ #
+
+ def test_injects_authorization_header_from_current_bearer_token(self):
+ self.vault_client.get_current_bearer_token.return_value = "the-current-token"
+ self.insert_api.with_raw_response.insert.return_value = FakeRawResponse([])
+
+ self.vault.insert(InsertRequest(values=[dict(values={"a": 1})], table="t1"))
+
+ _, kwargs = self.insert_api.with_raw_response.insert.call_args
+ headers = kwargs["request_options"]["additional_headers"]
+ self.assertEqual(headers.get("Authorization"), "Bearer the-current-token")
+
+ def test_no_authorization_header_when_no_token_available(self):
+ self.vault_client.get_current_bearer_token.return_value = None
+ self.insert_api.with_raw_response.insert.return_value = FakeRawResponse([])
+
+ self.vault.insert(InsertRequest(values=[dict(values={"a": 1})], table="t1"))
+
+ _, kwargs = self.insert_api.with_raw_response.insert.call_args
+ headers = kwargs["request_options"]["additional_headers"]
+ self.assertNotIn("Authorization", headers)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/flowvault/tests/vault/data/__init__.py b/flowvault/tests/vault/data/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/flowvault/tests/vault/data/test_data_classes.py b/flowvault/tests/vault/data/test_data_classes.py
new file mode 100644
index 00000000..33b274ca
--- /dev/null
+++ b/flowvault/tests/vault/data/test_data_classes.py
@@ -0,0 +1,57 @@
+import unittest
+
+from common.vault.data import BaseInsertRequest, BaseInsertResponse
+from skyflow_flowvault.utils.enums import UpsertType
+from skyflow_flowvault.vault.data import InsertRequest, InsertResponse
+
+
+class TestInsertRequest(unittest.TestCase):
+ def test_is_a_base_insert_request(self):
+ request = InsertRequest(values=[{"values": {"a": 1}}], table="t1")
+ self.assertIsInstance(request, BaseInsertRequest)
+ self.assertEqual(request.table, "t1")
+
+ def test_records_are_plain_dicts_supporting_per_record_overrides(self):
+ upsert = {"update_type": UpsertType.REPLACE, "unique_columns": ["a"]}
+ record = {"values": {"a": 1}, "table": "t2", "upsert": upsert}
+ request = InsertRequest(values=[record])
+ self.assertEqual(request.values[0]["values"], {"a": 1})
+ self.assertEqual(request.values[0]["table"], "t2")
+ self.assertIs(request.values[0]["upsert"], upsert)
+
+ def test_table_and_upsert_are_optional_defaults(self):
+ request = InsertRequest(values=[{"values": {"a": 1}}])
+ self.assertIsNone(request.table)
+ self.assertIsNone(request.upsert)
+
+ def test_no_v2_only_fields_exist(self):
+ request = InsertRequest(values=[{"values": {"a": 1}}])
+ for legacy_field in ("tokens", "homogeneous", "continue_on_error", "token_mode", "return_tokens"):
+ self.assertFalse(hasattr(request, legacy_field), f"v3 InsertRequest should not have '{legacy_field}'")
+
+
+class TestInsertResponse(unittest.TestCase):
+ """Shared shape with PDB's InsertResponse -- inserted_fields/errors, each entry tagged
+ request_index -- plain dicts/list-of-dicts, not custom classes."""
+
+ def test_shape(self):
+ inserted_fields = [{"request_index": 0, "skyflow_id": "id1"}]
+ response = InsertResponse(inserted_fields=inserted_fields, errors=[])
+
+ self.assertIs(response.inserted_fields, inserted_fields)
+ self.assertEqual(response.errors, [])
+
+ def test_is_a_base_insert_response(self):
+ response = InsertResponse(inserted_fields=[], errors=None)
+ self.assertIsInstance(response, BaseInsertResponse)
+
+ def test_repr_does_not_raise(self):
+ response = InsertResponse(
+ inserted_fields=[],
+ errors=[{"request_index": 0, "error": "boom", "code": 500, "request_id": None}],
+ )
+ self.assertIn("InsertResponse", repr(response))
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/ruff.toml b/ruff.toml
index aea6cce7..103d80d3 100644
--- a/ruff.toml
+++ b/ruff.toml
@@ -1,7 +1,7 @@
# ruff.toml
exclude = [
- "skyflow/generated",
+ "generated",
".git",
".ruff_cache",
".venv",
diff --git a/setup.py b/setup.py
deleted file mode 100644
index cfed2f0c..00000000
--- a/setup.py
+++ /dev/null
@@ -1,48 +0,0 @@
-'''
- Copyright (c) 2022 Skyflow, Inc.
-'''
-from setuptools import setup, find_packages
-import sys
-
-
-if sys.version_info < (3, 9):
- raise RuntimeError("skyflow requires Python 3.9+")
-current_version = '2.1.2'
-
-with open('README.md', 'r', encoding='utf-8') as f:
- long_description = f.read()
-
-setup(
- name='skyflow',
- version=current_version,
- author='Skyflow',
- author_email='service-ops@skyflow.com',
- packages=find_packages(where='.', exclude=['test*']),
- # Ship PEP 561 markers so type checkers (mypy/pyright) see the SDK's types.
- package_data={
- 'skyflow': ['py.typed'],
- 'skyflow.generated.rest': ['py.typed'],
- },
- url='https://github.com/skyflowapi/skyflow-python/',
- license='LICENSE',
- description='Skyflow SDK for the Python programming language',
- long_description=long_description,
- long_description_content_type='text/markdown',
- install_requires=[
- 'pydantic >= 2.0.0',
- 'typing-extensions >= 4.0.0',
- 'PyJWT >= 2.12, < 3',
- 'requests >= 2.28.0',
- 'cryptography >= 44.0.2',
- 'httpx >= 0.21.2',
- 'python-dotenv >= 1.1.0, < 2',
- ],
- extras_require={
- 'dev': [
- 'codespell >= 2.4.1',
- 'ruff >= 0.9.0',
- 'pre-commit >= 4.3.0',
- ]
- },
- python_requires=">=3.9",
-)
diff --git a/skyflow/client/skyflow.py b/skyflow/client/skyflow.py
deleted file mode 100644
index ebd5ef7d..00000000
--- a/skyflow/client/skyflow.py
+++ /dev/null
@@ -1,253 +0,0 @@
-from collections import OrderedDict
-from skyflow import LogLevel
-from skyflow.error import SkyflowError
-from skyflow.utils import SkyflowMessages
-from skyflow.utils.logger import log_info, log_warn, set_active_log_level, Logger
-from skyflow.utils.constants import OptionField
-from skyflow.utils.validations import validate_vault_config, validate_connection_config, validate_update_vault_config, \
- validate_update_connection_config, validate_credentials, validate_log_level
-from skyflow.vault.client.client import VaultClient
-from skyflow.vault.controller import Vault
-from skyflow.vault.controller import Connection
-from skyflow.vault.controller import Detect
-
-class Skyflow:
- def __init__(self, builder):
- self.__builder = builder
- log_info(SkyflowMessages.Info.CLIENT_INITIALIZED.value, self.__builder.get_logger())
-
- @staticmethod
- def builder():
- return Skyflow.Builder()
-
- def add_vault_config(self, config):
- self.__builder._Builder__add_vault_config(config)
- return self
-
- def remove_vault_config(self, vault_id):
- self.__builder.remove_vault_config(vault_id)
-
- def update_vault_config(self,config):
- self.__builder.update_vault_config(config)
-
- def get_vault_config(self, vault_id):
- return self.__builder.get_vault_config(vault_id).get(OptionField.VAULT_CLIENT).get_config()
-
- def add_connection_config(self, config):
- self.__builder._Builder__add_connection_config(config)
- return self
-
- def remove_connection_config(self, connection_id):
- self.__builder.remove_connection_config(connection_id)
- return self
-
- def update_connection_config(self, config):
- self.__builder.update_connection_config(config)
- return self
-
- def get_connection_config(self, connection_id):
- return self.__builder.get_connection_config(connection_id).get(OptionField.VAULT_CLIENT).get_config()
-
- def add_skyflow_credentials(self, credentials):
- self.__builder._Builder__add_skyflow_credentials(credentials)
- return self
-
- def update_skyflow_credentials(self, credentials):
- self.__builder._Builder__add_skyflow_credentials(credentials)
-
- def set_log_level(self, log_level):
- self.__builder._Builder__set_log_level(log_level)
- return self
-
- def update_log_level(self, log_level):
- """.. deprecated:: Use set_log_level() instead. Will be removed in a future release."""
- log_warn(SkyflowMessages.Warning.UPDATE_LOG_LEVEL_DEPRECATED.value)
- return self.set_log_level(log_level)
-
- def get_log_level(self):
- return self.__builder._Builder__log_level
-
- def vault(self, vault_id = None) -> Vault:
- vault_config = self.__builder.get_vault_config(vault_id)
- return vault_config.get(OptionField.VAULT_CONTROLLER)
-
- def connection(self, connection_id = None) -> Connection:
- connection_config = self.__builder.get_connection_config(connection_id)
- return connection_config.get(OptionField.CONTROLLER)
-
- def detect(self, vault_id = None) -> Detect:
- vault_config = self.__builder.get_vault_config(vault_id)
- return vault_config.get(OptionField.DETECT_CONTROLLER)
-
- class Builder:
- def __init__(self):
- self.__vault_configs = OrderedDict()
- self.__vault_list = list()
- self.__connection_configs = OrderedDict()
- self.__connection_list = list()
- self.__skyflow_credentials = None
- self.__log_level = LogLevel.ERROR
- self.__logger = Logger(LogLevel.ERROR)
-
- def add_vault_config(self, config):
- vault_id = config.get(OptionField.VAULT_ID)
- if not isinstance(vault_id, str) or not vault_id:
- raise SkyflowError(
- SkyflowMessages.Error.INVALID_VAULT_ID.value,
- SkyflowMessages.ErrorCodes.INVALID_INPUT.value
- )
- if vault_id in [vault.get(OptionField.VAULT_ID) for vault in self.__vault_list]:
- log_info(SkyflowMessages.Info.VAULT_CONFIG_EXISTS.value.format(vault_id), self.__logger)
- raise SkyflowError(
- SkyflowMessages.Error.VAULT_ID_ALREADY_EXISTS.value.format(vault_id),
- SkyflowMessages.ErrorCodes.INVALID_INPUT.value
- )
-
- self.__vault_list.append(config)
- return self
-
- def remove_vault_config(self, vault_id):
- if vault_id in self.__vault_configs.keys():
- self.__vault_configs.pop(vault_id)
- else:
- raise SkyflowError(SkyflowMessages.Error.INVALID_VAULT_ID.value,
- SkyflowMessages.ErrorCodes.INVALID_INPUT.value)
-
- def update_vault_config(self, config):
- validate_update_vault_config(self.__logger, config)
- vault_id = config.get(OptionField.VAULT_ID)
- if vault_id not in self.__vault_configs:
- raise SkyflowError(SkyflowMessages.Error.VAULT_ID_NOT_IN_CONFIG_LIST.value.format(vault_id), SkyflowMessages.ErrorCodes.INVALID_INPUT.value)
- vault_config = self.__vault_configs[vault_id]
- vault_config.get(OptionField.VAULT_CLIENT).update_config(config)
-
- def get_vault_config(self, vault_id):
- if vault_id is None:
- if self.__vault_configs:
- return next(iter(self.__vault_configs.values()))
- raise SkyflowError(SkyflowMessages.Error.EMPTY_VAULT_CONFIGS.value, SkyflowMessages.ErrorCodes.INVALID_INPUT.value)
-
- if vault_id in self.__vault_configs:
- return self.__vault_configs.get(vault_id)
- log_info(SkyflowMessages.Info.VAULT_CONFIG_DOES_NOT_EXIST.value.format(vault_id), self.__logger)
- raise SkyflowError(SkyflowMessages.Error.VAULT_ID_NOT_IN_CONFIG_LIST.value.format(vault_id), SkyflowMessages.ErrorCodes.INVALID_INPUT.value)
-
-
- def add_connection_config(self, config):
- connection_id = config.get(OptionField.CONNECTION_ID)
- if not isinstance(connection_id, str) or not connection_id:
- raise SkyflowError(
- SkyflowMessages.Error.INVALID_CONNECTION_ID.value,
- SkyflowMessages.ErrorCodes.INVALID_INPUT.value
- )
- if connection_id in [connection.get(OptionField.CONNECTION_ID) for connection in self.__connection_list]:
- log_info(SkyflowMessages.Info.CONNECTION_CONFIG_EXISTS.value.format(connection_id), self.__logger)
- raise SkyflowError(
- SkyflowMessages.Error.CONNECTION_ID_ALREADY_EXISTS.value.format(connection_id),
- SkyflowMessages.ErrorCodes.INVALID_INPUT.value
- )
- self.__connection_list.append(config)
- return self
-
- def remove_connection_config(self, connection_id):
- if connection_id in self.__connection_configs.keys():
- self.__connection_configs.pop(connection_id)
- else:
- raise SkyflowError(SkyflowMessages.Error.INVALID_CONNECTION_ID.value,
- SkyflowMessages.ErrorCodes.INVALID_INPUT.value)
-
- def update_connection_config(self, config):
- validate_update_connection_config(self.__logger, config)
- connection_id = config[OptionField.CONNECTION_ID]
- if connection_id not in self.__connection_configs:
- raise SkyflowError(SkyflowMessages.Error.CONNECTION_ID_NOT_IN_CONFIG_LIST.value.format(connection_id), SkyflowMessages.ErrorCodes.INVALID_INPUT.value)
- connection_config = self.__connection_configs[connection_id]
- connection_config.get(OptionField.VAULT_CLIENT).update_config(config)
-
- def get_connection_config(self, connection_id):
- if connection_id is None:
- if self.__connection_configs:
- return next(iter(self.__connection_configs.values()))
-
- raise SkyflowError(SkyflowMessages.Error.EMPTY_CONNECTION_CONFIGS.value, SkyflowMessages.ErrorCodes.INVALID_INPUT.value)
-
- if connection_id in self.__connection_configs:
- return self.__connection_configs.get(connection_id)
- log_info(SkyflowMessages.Info.CONNECTION_CONFIG_DOES_NOT_EXIST.value.format(connection_id), self.__logger)
- raise SkyflowError(SkyflowMessages.Error.CONNECTION_ID_NOT_IN_CONFIG_LIST.value.format(connection_id), SkyflowMessages.ErrorCodes.INVALID_INPUT.value)
-
-
- def add_skyflow_credentials(self, credentials):
- self.__skyflow_credentials = credentials
- return self
-
- def set_log_level(self, log_level):
- self.__log_level = log_level
- return self
-
- def get_logger(self):
- return self.__logger
-
- def __add_vault_config(self, config):
- validate_vault_config(self.__logger, config)
- vault_id = config.get(OptionField.VAULT_ID)
- vault_client = VaultClient(config)
- self.__vault_configs[vault_id] = {
- OptionField.VAULT_CLIENT: vault_client,
- OptionField.VAULT_CONTROLLER: Vault(vault_client),
- OptionField.DETECT_CONTROLLER: Detect(vault_client)
- }
- log_info(SkyflowMessages.Info.VAULT_CONTROLLER_INITIALIZED.value.format(config.get(OptionField.VAULT_ID)), self.__logger)
- log_info(SkyflowMessages.Info.DETECT_CONTROLLER_INITIALIZED.value.format(config.get(OptionField.VAULT_ID)), self.__logger)
-
- def __add_connection_config(self, config):
- validate_connection_config(self.__logger, config)
- connection_id = config.get(OptionField.CONNECTION_ID)
- vault_client = VaultClient(config)
- self.__connection_configs[connection_id] = {
- OptionField.VAULT_CLIENT: vault_client,
- OptionField.CONTROLLER: Connection(vault_client)
- }
- log_info(SkyflowMessages.Info.CONNECTION_CONTROLLER_INITIALIZED.value.format(config.get(OptionField.CONNECTION_ID)), self.__logger)
-
- def __update_vault_client_logger(self, log_level, logger):
- for vault_id, vault_config in self.__vault_configs.items():
- vault_config.get(OptionField.VAULT_CLIENT).set_logger(log_level,logger)
-
- for connection_id, connection_config in self.__connection_configs.items():
- connection_config.get(OptionField.VAULT_CLIENT).set_logger(log_level,logger)
-
- def __set_log_level(self, log_level):
- validate_log_level(self.__logger, log_level)
- self.__log_level = log_level
- self.__logger.set_log_level(log_level)
- set_active_log_level(log_level)
- self.__update_vault_client_logger(log_level, self.__logger)
- log_info(SkyflowMessages.Info.LOGGER_SETUP_DONE.value, self.__logger)
- log_info(SkyflowMessages.Info.CURRENT_LOG_LEVEL.value.format(self.__log_level), self.__logger)
-
- def __add_skyflow_credentials(self, credentials):
- if credentials is not None:
- self.__skyflow_credentials = credentials
- validate_credentials(self.__logger, credentials)
- for vault_id, vault_config in self.__vault_configs.items():
- vault_config.get(OptionField.VAULT_CLIENT).set_common_skyflow_credentials(credentials)
-
- for connection_id, connection_config in self.__connection_configs.items():
- connection_config.get(OptionField.VAULT_CLIENT).set_common_skyflow_credentials(self.__skyflow_credentials)
- def build(self):
- validate_log_level(self.__logger, self.__log_level)
- self.__logger.set_log_level(self.__log_level)
- set_active_log_level(self.__log_level)
-
- for config in self.__vault_list:
- self.__add_vault_config(config)
-
- for config in self.__connection_list:
- self.__add_connection_config(config)
-
- self.__update_vault_client_logger(self.__log_level, self.__logger)
-
- self.__add_skyflow_credentials(self.__skyflow_credentials)
-
- return Skyflow(self)
diff --git a/skyflow/error/__init__.py b/skyflow/error/__init__.py
deleted file mode 100644
index 305c7966..00000000
--- a/skyflow/error/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from ._skyflow_error import SkyflowError
\ No newline at end of file
diff --git a/skyflow/vault/client/client.py b/skyflow/vault/client/client.py
deleted file mode 100644
index 8023646c..00000000
--- a/skyflow/vault/client/client.py
+++ /dev/null
@@ -1,119 +0,0 @@
-from skyflow.error import SkyflowError
-from skyflow.generated.rest.client import Skyflow
-from skyflow.service_account import generate_bearer_token, generate_bearer_token_from_creds, is_expired
-from skyflow.utils import get_vault_url, get_credentials, SkyflowMessages
-from skyflow.utils.logger import log_info
-from skyflow.utils.constants import OptionField, CredentialField, ConfigField
-
-
-class VaultClient:
- def __init__(self, config):
- self.__config = config
- self.__common_skyflow_credentials = None
- self.__log_level = None
- self.__client_configuration = None
- self.__api_client = None
- self.__logger = None
- self.__is_config_updated = False
- self.__bearer_token = None
- self.__credentials = None
- self.__vault_url = None
- self.__is_static_token = None
-
- def set_common_skyflow_credentials(self, credentials):
- self.__common_skyflow_credentials = credentials
-
- def set_logger(self, log_level, logger):
- self.__log_level = log_level
- self.__logger = logger
-
- def initialize_client_configuration(self):
- if self.__api_client is not None and not self.__is_config_updated:
- if self.__is_static_token:
- return
- if self.__bearer_token is not None and not is_expired(self.__bearer_token):
- return
-
- needs_reinit = self.__api_client is None or self.__is_config_updated
- if needs_reinit:
- self.__credentials = get_credentials(self.__config.get(ConfigField.CREDENTIALS), self.__common_skyflow_credentials, logger=self.__logger)
- self.__vault_url = get_vault_url(self.__config.get(ConfigField.CLUSTER_ID),
- self.__config.get(ConfigField.ENV),
- self.__config.get(ConfigField.VAULT_ID),
- logger=self.__logger)
- self.__is_static_token = CredentialField.TOKEN in self.__credentials or CredentialField.API_KEY in self.__credentials
- bearer_token = self.get_bearer_token(self.__credentials)
- if needs_reinit:
- self.initialize_api_client(self.__vault_url, bearer_token)
-
- def initialize_api_client(self, vault_url, bearer_token):
- token_provider = lambda: self.__bearer_token if self.__bearer_token is not None else bearer_token # noqa: E731
- self.__api_client = Skyflow(base_url=vault_url, token=token_provider)
-
- def get_records_api(self):
- return self.__api_client.records
-
- def get_tokens_api(self):
- return self.__api_client.tokens
-
- def get_query_api(self):
- return self.__api_client.query
-
- def get_detect_text_api(self):
- return self.__api_client.strings
-
- def get_detect_file_api(self):
- return self.__api_client.files
-
- def get_vault_id(self):
- return self.__config.get(ConfigField.VAULT_ID)
-
- def get_bearer_token(self, credentials):
- if CredentialField.API_KEY in credentials:
- return credentials.get(CredentialField.API_KEY)
- elif CredentialField.TOKEN in credentials:
- return credentials.get(CredentialField.TOKEN)
-
- options = {
- OptionField.ROLE_IDS: self.__config.get(OptionField.ROLES),
- OptionField.CTX: self.__config.get(OptionField.CTX)
- }
- if CredentialField.TOKEN_URI_OPTION in credentials and credentials.get(CredentialField.TOKEN_URI_OPTION):
- options[CredentialField.TOKEN_URI_OPTION] = credentials.get(CredentialField.TOKEN_URI_OPTION)
-
- if self.__bearer_token is None or self.__is_config_updated or is_expired(self.__bearer_token):
- if CredentialField.PATH in credentials:
- self.__bearer_token, _ = generate_bearer_token(
- credentials.get(CredentialField.PATH),
- options,
- self.__logger
- )
- else:
- credentials_string = credentials.get(CredentialField.CREDENTIALS_STRING)
- log_info(SkyflowMessages.Info.GENERATE_BEARER_TOKEN_FROM_CREDENTIALS_STRING_TRIGGERED.value, self.__logger)
- self.__bearer_token, _ = generate_bearer_token_from_creds(
- credentials_string,
- options,
- self.__logger
- )
- self.__is_config_updated = False
- else:
- log_info(SkyflowMessages.Info.REUSE_BEARER_TOKEN.value, self.__logger)
-
- return self.__bearer_token
-
- def update_config(self, config):
- self.__config.update(config)
- self.__is_config_updated = True
-
- def get_config(self):
- return self.__config
-
- def get_common_skyflow_credentials(self):
- return self.__common_skyflow_credentials
-
- def get_log_level(self):
- return self.__log_level
-
- def get_logger(self):
- return self.__logger
\ No newline at end of file
diff --git a/skyflow/vault/controller/__init__.py b/skyflow/vault/controller/__init__.py
deleted file mode 100644
index 681153c0..00000000
--- a/skyflow/vault/controller/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from ._vault import Vault
-from ._connections import Connection
-from ._detect import Detect
\ No newline at end of file
diff --git a/skyflow/vault/data/_insert_request.py b/skyflow/vault/data/_insert_request.py
deleted file mode 100644
index 909edd88..00000000
--- a/skyflow/vault/data/_insert_request.py
+++ /dev/null
@@ -1,21 +0,0 @@
-from skyflow.utils.enums import TokenMode
-
-class InsertRequest:
- def __init__(self,
- table,
- values,
- tokens = None,
- upsert = None,
- homogeneous = False,
- token_mode = TokenMode.DISABLE,
- return_tokens = True,
- continue_on_error = False):
- self.table = table
- self.values = values
- self.tokens = tokens
- self.upsert = upsert
- self.homogeneous = homogeneous
- self.token_mode = token_mode
- self.return_tokens = return_tokens
- self.continue_on_error = continue_on_error
-
diff --git a/skyflow/vault/data/_insert_response.py b/skyflow/vault/data/_insert_response.py
deleted file mode 100644
index 0c7c777f..00000000
--- a/skyflow/vault/data/_insert_response.py
+++ /dev/null
@@ -1,10 +0,0 @@
-class InsertResponse:
- def __init__(self, inserted_fields = None, errors=None):
- self.inserted_fields = inserted_fields
- self.errors = errors
-
- def __repr__(self):
- return f"InsertResponse(inserted_fields={self.inserted_fields}, errors={self.errors})"
-
- def __str__(self):
- return self.__repr__()
diff --git a/tests/contract/__init__.py b/tests/contract/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/contract/_adapter_loader.py b/tests/contract/_adapter_loader.py
new file mode 100644
index 00000000..2534a174
--- /dev/null
+++ b/tests/contract/_adapter_loader.py
@@ -0,0 +1,26 @@
+"""Selects the v2 or v3 (flowvault) contract adapter based on SKYFLOW_TEST_VARIANT.
+Plain-Python equivalent of a pytest conftest.py fixture -- this repo's test runner is plain
+unittest (see each variant's tests/), so variant selection happens via import rather than a
+fixture.
+
+Usage (run once per variant, in that variant's own installed/PYTHONPATH environment -- v2's
+skyflow and flowvault's skyflow_flowvault can never coexist in one process):
+
+ SKYFLOW_TEST_VARIANT=v2 PYTHONPATH=.:v2 python -m unittest discover -s tests/contract -t .
+ SKYFLOW_TEST_VARIANT=v3 PYTHONPATH=.:flowvault python -m unittest discover -s tests/contract -t .
+"""
+import os
+
+_VARIANT = os.environ.get("SKYFLOW_TEST_VARIANT")
+
+if _VARIANT == "v2":
+ from tests.contract.adapters import v2_adapter as adapter
+elif _VARIANT == "v3":
+ from tests.contract.adapters import v3_adapter as adapter
+else:
+ raise RuntimeError(
+ "SKYFLOW_TEST_VARIANT must be set to 'v2' or 'v3' before running tests/contract/ "
+ "(e.g. SKYFLOW_TEST_VARIANT=v2 PYTHONPATH=.:v2 python -m unittest discover -s tests/contract -t .)"
+ )
+
+VARIANT = _VARIANT
diff --git a/tests/contract/adapters/__init__.py b/tests/contract/adapters/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/contract/adapters/v2_adapter.py b/tests/contract/adapters/v2_adapter.py
new file mode 100644
index 00000000..ff581328
--- /dev/null
+++ b/tests/contract/adapters/v2_adapter.py
@@ -0,0 +1,52 @@
+"""Contract adapter for v2. Constructs a Vault with its underlying generated API call mocked
+out, so the contract suite can assert on SDK-level behavior (call counts, response shape)
+without needing real network access or credentials. Import this module only from a v2-installed
+environment (`SKYFLOW_TEST_VARIANT=v2`) -- v2.skyflow and v3.skyflow cannot coexist in one
+process."""
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+from skyflow.vault.client.client import VaultClient
+from skyflow.vault.controller import Vault
+from skyflow.vault.data import InsertRequest
+
+# Uses the public `Vault` alias deliberately (not VaultController) -- this adapter simulates
+# an external consumer, and Vault is the name they'd actually import.
+
+
+def build_vault():
+ config = {
+ "vault_id": "contract_vault",
+ "cluster_id": "contract_cluster",
+ "env": "PROD",
+ "credentials": {"token": "contract_static_token"},
+ }
+ vault_client = VaultClient(config)
+ vault_client.initialize_client_configuration = MagicMock() # skip real credential/URL resolution
+ records_api = MagicMock()
+ vault_client.get_records_api = MagicMock(return_value=records_api)
+ vault = Vault(vault_client)
+ return vault, records_api
+
+
+def build_insert_request(n):
+ return InsertRequest(table="contract_table", values=[{"field": f"value{i}"} for i in range(n)])
+
+
+def call_insert(vault, records_api, request):
+ fake_records = [SimpleNamespace(skyflow_id=f"id{i}", tokens=None) for i in range(len(request.values))]
+ fake_response = SimpleNamespace(data=SimpleNamespace(records=fake_records), headers={})
+ records_api.with_raw_response.record_service_insert_record.return_value = fake_response
+
+ response = vault.insert(request)
+ call_count = records_api.with_raw_response.record_service_insert_record.call_count
+ return response, call_count
+
+
+# v2's InsertResponse is the shared shape (inserted_fields/errors) both variants now use.
+def count_successes(response):
+ return len(response.inserted_fields)
+
+
+def count_errors(response):
+ return len(response.errors) if response.errors else 0
diff --git a/tests/contract/adapters/v3_adapter.py b/tests/contract/adapters/v3_adapter.py
new file mode 100644
index 00000000..d5a7ceb3
--- /dev/null
+++ b/tests/contract/adapters/v3_adapter.py
@@ -0,0 +1,52 @@
+"""Contract adapter for v3 (flowvault). See v2_adapter.py for the shared design note -- import
+this module only from a flowvault-installed environment (`SKYFLOW_TEST_VARIANT=v3`)."""
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+from skyflow_flowvault.vault.client.client import VaultClient
+from skyflow_flowvault.vault.controller import VaultController
+from skyflow_flowvault.vault.data import InsertRequest
+
+
+def build_vault():
+ config = {
+ "vault_id": "contract_vault",
+ "cluster_id": "contract_cluster",
+ "env": "PROD",
+ "credentials": {"token": "contract_static_token"},
+ }
+ vault_client = VaultClient(config)
+ vault_client.initialize_client_configuration = MagicMock() # skip real credential/URL resolution
+ insert_api = MagicMock()
+ vault_client.get_insert_api = MagicMock(return_value=insert_api)
+ vault = VaultController(vault_client)
+ return vault, insert_api
+
+
+def build_insert_request(n):
+ return InsertRequest(table="contract_table", values=[dict(values={"field": f"value{i}"}) for i in range(n)])
+
+
+def call_insert(vault, insert_api, request):
+ def fake_insert(**kwargs):
+ records = [
+ SimpleNamespace(skyflow_id=f"id{i}", tokens=None, data=None, error=None, http_code=None, table_name=None)
+ for i in range(len(kwargs["records"]))
+ ]
+ return SimpleNamespace(data=SimpleNamespace(records=records), headers={})
+
+ insert_api.with_raw_response.insert.side_effect = fake_insert
+ response = vault.insert(request)
+ call_count = insert_api.with_raw_response.insert.call_count
+ return response, call_count
+
+
+# v3's InsertResponse now shares the exact same shape as v2's (inserted_fields/errors, each
+# entry tagged request_index) -- kept as separate accessor functions per adapter anyway, since
+# the contract module intentionally treats each variant's response as opaque.
+def count_successes(response):
+ return len(response.inserted_fields)
+
+
+def count_errors(response):
+ return len(response.errors) if response.errors else 0
diff --git a/tests/contract/test_insert_contract.py b/tests/contract/test_insert_contract.py
new file mode 100644
index 00000000..69406407
--- /dev/null
+++ b/tests/contract/test_insert_contract.py
@@ -0,0 +1,48 @@
+"""Shared insert() contract, asserted identically against both variants -- authored once, run
+twice (see _adapter_loader.py). If this file needs variant-specific branching, that's a signal
+the abstraction leaked and the adapter interface needs to grow, not this file.
+"""
+import unittest
+
+from tests.contract._adapter_loader import VARIANT, adapter
+
+
+class TestInsertContract(unittest.TestCase):
+ def test_vault_exposes_insert_with_a_single_request_argument(self):
+ vault, _ = adapter.build_vault()
+ self.assertTrue(hasattr(vault, "insert"), f"{VARIANT}'s Vault must expose insert()")
+ self.assertTrue(callable(vault.insert))
+
+ def test_insert_response_reports_correct_counts(self):
+ """v2 and v3 now share the exact same InsertResponse shape (inserted_fields/errors,
+ each entry tagged request_index); the contract accesses it via the adapter's
+ count_successes/count_errors functions rather than assuming the shape directly."""
+ vault, api = adapter.build_vault()
+ request = adapter.build_insert_request(1)
+
+ response, _ = adapter.call_insert(vault, api, request)
+
+ self.assertEqual(adapter.count_successes(response), 1)
+ self.assertEqual(adapter.count_errors(response), 0)
+
+ def test_insert_of_many_records_returns_one_field_per_record(self):
+ vault, api = adapter.build_vault()
+ request = adapter.build_insert_request(7)
+
+ response, _ = adapter.call_insert(vault, api, request)
+
+ self.assertEqual(adapter.count_successes(response), 7)
+
+ def test_insert_always_makes_exactly_one_api_call_regardless_of_record_count(self):
+ """Neither variant batches -- every insert(), no matter how many records, must reach
+ the underlying API exactly once."""
+ vault, api = adapter.build_vault()
+ request = adapter.build_insert_request(3)
+
+ _, call_count = adapter.call_insert(vault, api, request)
+
+ self.assertEqual(call_count, 1, f"{VARIANT} must always call the underlying API exactly once")
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/contract/test_typecheck_contract.py b/tests/contract/test_typecheck_contract.py
new file mode 100644
index 00000000..df2cbf92
--- /dev/null
+++ b/tests/contract/test_typecheck_contract.py
@@ -0,0 +1,54 @@
+"""Only meaningful under SKYFLOW_TEST_VARIANT=v3 -- shells out to mypy (or pyright, if mypy
+isn't installed) against the fixture and asserts it fails with diagnostics naming each v2-only
+field. Skipped entirely under v2 (those kwargs are valid there).
+
+NOT independently verified in this environment: neither mypy nor pyright is installed in the
+sandbox this was authored in, so this file has been reviewed for correctness but its actual
+pass/fail behavior against a real type checker has not been observed. Install one of them (`pip
+install mypy` or `pip install pyright`) before relying on this as a passing gate.
+"""
+import shutil
+import subprocess
+import sys
+import unittest
+from pathlib import Path
+
+from tests.contract._adapter_loader import VARIANT
+
+FIXTURE = Path(__file__).parent / "typecheck_fixtures" / "v2_only_insert_kwargs_should_fail_under_v3.py"
+V2_ONLY_FIELDS = ["homogeneous", "continue_on_error", "token_mode", "return_tokens"]
+
+
+@unittest.skipUnless(VARIANT == "v3", "v2-only insert kwargs are valid under v2; nothing to type-check there")
+class TestV3RejectsV2OnlyInsertFields(unittest.TestCase):
+ def test_v2_only_kwargs_are_flagged_as_type_errors_under_v3(self):
+ if shutil.which("mypy") or _module_available("mypy"):
+ result = subprocess.run(
+ [sys.executable, "-m", "mypy", "--no-error-summary", str(FIXTURE)],
+ capture_output=True, text=True,
+ )
+ elif shutil.which("pyright") or _module_available("pyright"):
+ result = subprocess.run(
+ [sys.executable, "-m", "pyright", str(FIXTURE)],
+ capture_output=True, text=True,
+ )
+ else:
+ self.skipTest("neither mypy nor pyright is installed")
+ return
+
+ diagnostics = result.stdout + result.stderr
+ self.assertNotEqual(result.returncode, 0, f"expected type errors, got a clean run:\n{diagnostics}")
+ for field in V2_ONLY_FIELDS:
+ self.assertIn(field, diagnostics, f"expected a diagnostic mentioning '{field}':\n{diagnostics}")
+
+
+def _module_available(name):
+ try:
+ __import__(name)
+ return True
+ except ImportError:
+ return False
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/contract/typecheck_fixtures/v2_only_insert_kwargs_should_fail_under_v3.py b/tests/contract/typecheck_fixtures/v2_only_insert_kwargs_should_fail_under_v3.py
new file mode 100644
index 00000000..84c91c69
--- /dev/null
+++ b/tests/contract/typecheck_fixtures/v2_only_insert_kwargs_should_fail_under_v3.py
@@ -0,0 +1,19 @@
+"""Type-check fixture, not a test file to be executed directly (see test_typecheck_contract.py).
+
+Every construction below uses a v2-only InsertRequest keyword argument that does not exist on
+flowvault's InsertRequest (records/table/upsert only -- see
+flowvault/skyflow_flowvault/vault/data/_insert_request.py). Run under mypy/pyright against a
+flowvault install, each of these lines must be flagged as a type error. If a future change to
+flowvault's InsertRequest ever silently grows one of these fields back, this fixture stops
+producing errors and test_typecheck_contract.py's assertion on it fails -- that's the point:
+it's a regression trip-wire, not a demonstration.
+
+No `# type: ignore` anywhere in this file -- the whole point is for the checker to actually emit
+diagnostics.
+"""
+from skyflow_flowvault.vault.data import InsertRequest
+
+InsertRequest(values=[], table="t1", homogeneous=True)
+InsertRequest(values=[], table="t1", continue_on_error=True)
+InsertRequest(values=[], table="t1", token_mode="ENABLE")
+InsertRequest(values=[], table="t1", return_tokens=False)
diff --git a/tests/vault/client/test__client.py b/tests/vault/client/test__client.py
deleted file mode 100644
index 4df508c7..00000000
--- a/tests/vault/client/test__client.py
+++ /dev/null
@@ -1,327 +0,0 @@
-import unittest
-from unittest.mock import patch, MagicMock
-
-from skyflow.error import SkyflowError
-from skyflow.utils import SkyflowMessages
-from skyflow.vault.client.client import VaultClient
-
-CONFIG = {
- "credentials": "some_credentials",
- "cluster_id": "test_cluster_id",
- "env": "test_env",
- "vault_id": "test_vault_id",
- "roles": ["role_id_1", "role_id_2"],
- "ctx": "context"
-}
-
-CREDENTIALS_WITH_API_KEY = {"api_key": "dummy_api_key"}
-CREDENTIALS_WITH_TOKEN = {"token": "dummy_static_token"}
-CREDENTIALS_WITH_PATH = {"path": "/some/path/credentials.json"}
-CREDENTIALS_WITH_STRING = {"credentials_string": '{"clientID": "x"}'}
-
-
-class TestVaultClient(unittest.TestCase):
- def setUp(self):
- self.vault_client = VaultClient(CONFIG)
-
- # ------------------------------------------------------------------ #
- # Basic setters / getters #
- # ------------------------------------------------------------------ #
-
- def test_set_common_skyflow_credentials(self):
- credentials = {"api_key": "dummy_api_key"}
- self.vault_client.set_common_skyflow_credentials(credentials)
- self.assertEqual(self.vault_client.get_common_skyflow_credentials(), credentials)
-
- def test_set_logger(self):
- mock_logger = MagicMock()
- self.vault_client.set_logger("INFO", mock_logger)
- self.assertEqual(self.vault_client.get_log_level(), "INFO")
- self.assertEqual(self.vault_client.get_logger(), mock_logger)
-
- def test_get_vault_id(self):
- self.assertEqual(self.vault_client.get_vault_id(), CONFIG["vault_id"])
-
- def test_get_config(self):
- self.assertEqual(self.vault_client.get_config(), CONFIG)
-
- def test_get_common_skyflow_credentials(self):
- credentials = {"api_key": "dummy_api_key"}
- self.vault_client.set_common_skyflow_credentials(credentials)
- self.assertEqual(self.vault_client.get_common_skyflow_credentials(), credentials)
-
- def test_get_log_level(self):
- self.vault_client.set_logger("DEBUG", MagicMock())
- self.assertEqual(self.vault_client.get_log_level(), "DEBUG")
-
- def test_get_logger(self):
- mock_logger = MagicMock()
- self.vault_client.set_logger("INFO", mock_logger)
- self.assertEqual(self.vault_client.get_logger(), mock_logger)
-
- # ------------------------------------------------------------------ #
- # initialize_client_configuration — first call (slow path) #
- # ------------------------------------------------------------------ #
-
- @patch("skyflow.vault.client.client.get_credentials")
- @patch("skyflow.vault.client.client.get_vault_url")
- @patch("skyflow.vault.client.client.VaultClient.initialize_api_client")
- def test_initialize_client_configuration_first_call(
- self, mock_init_api_client, mock_get_vault_url, mock_get_credentials
- ):
- mock_get_credentials.return_value = CREDENTIALS_WITH_API_KEY
- mock_get_vault_url.return_value = "https://test-vault-url.com"
-
- self.vault_client.initialize_client_configuration()
-
- mock_get_credentials.assert_called_once_with(
- CONFIG["credentials"], None, logger=None
- )
- mock_get_vault_url.assert_called_once_with(
- CONFIG["cluster_id"], CONFIG["env"], CONFIG["vault_id"], logger=None
- )
- mock_init_api_client.assert_called_once()
-
- # ------------------------------------------------------------------ #
- # initialize_client_configuration — fast path (static token) #
- # ------------------------------------------------------------------ #
-
- @patch("skyflow.vault.client.client.get_credentials")
- @patch("skyflow.vault.client.client.get_vault_url")
- @patch("skyflow.vault.client.client.VaultClient.initialize_api_client")
- def test_initialize_client_configuration_fast_path_api_key(
- self, mock_init_api_client, mock_get_vault_url, mock_get_credentials
- ):
- """Once initialized with api_key, subsequent calls skip all work."""
- mock_get_credentials.return_value = CREDENTIALS_WITH_API_KEY
- mock_get_vault_url.return_value = "https://test-vault-url.com"
- # Side-effect simulates initialize_api_client actually setting __api_client
- mock_init_api_client.side_effect = lambda *_: setattr(
- self.vault_client, "_VaultClient__api_client", MagicMock()
- )
-
- self.vault_client.initialize_client_configuration() # first call — slow path
- mock_get_credentials.reset_mock()
- mock_get_vault_url.reset_mock()
- mock_init_api_client.reset_mock()
-
- self.vault_client.initialize_client_configuration() # second call — fast path
-
- mock_get_credentials.assert_not_called()
- mock_get_vault_url.assert_not_called()
- mock_init_api_client.assert_not_called()
-
- @patch("skyflow.vault.client.client.get_credentials")
- @patch("skyflow.vault.client.client.get_vault_url")
- @patch("skyflow.vault.client.client.VaultClient.initialize_api_client")
- def test_initialize_client_configuration_fast_path_static_token(
- self, mock_init_api_client, mock_get_vault_url, mock_get_credentials
- ):
- """Once initialized with a static token, subsequent calls skip all work."""
- mock_get_credentials.return_value = CREDENTIALS_WITH_TOKEN
- mock_get_vault_url.return_value = "https://test-vault-url.com"
- mock_init_api_client.side_effect = lambda *_: setattr(
- self.vault_client, "_VaultClient__api_client", MagicMock()
- )
-
- self.vault_client.initialize_client_configuration()
- mock_get_credentials.reset_mock()
- mock_get_vault_url.reset_mock()
- mock_init_api_client.reset_mock()
-
- self.vault_client.initialize_client_configuration()
-
- mock_get_credentials.assert_not_called()
- mock_get_vault_url.assert_not_called()
- mock_init_api_client.assert_not_called()
-
- # ------------------------------------------------------------------ #
- # initialize_client_configuration — fast path (service account) #
- # ------------------------------------------------------------------ #
-
- @patch("skyflow.vault.client.client.is_expired", return_value=False)
- @patch("skyflow.vault.client.client.get_credentials")
- @patch("skyflow.vault.client.client.get_vault_url")
- @patch("skyflow.vault.client.client.VaultClient.initialize_api_client")
- def test_initialize_client_configuration_fast_path_valid_sa_token(
- self, mock_init_api_client, mock_get_vault_url, mock_get_credentials, mock_is_expired
- ):
- """Service account with a still-valid token skips get_bearer_token entirely."""
- mock_get_credentials.return_value = CREDENTIALS_WITH_PATH
- mock_get_vault_url.return_value = "https://test-vault-url.com"
-
- # Seed the cached bearer token as if first call already ran
- self.vault_client._VaultClient__api_client = MagicMock()
- self.vault_client._VaultClient__is_static_token = False
- self.vault_client._VaultClient__bearer_token = "cached_sa_token"
- self.vault_client._VaultClient__credentials = CREDENTIALS_WITH_PATH
-
- self.vault_client.initialize_client_configuration()
-
- mock_get_credentials.assert_not_called()
- mock_get_vault_url.assert_not_called()
- mock_init_api_client.assert_not_called()
-
- # ------------------------------------------------------------------ #
- # initialize_client_configuration — token expiry (no client reinit) #
- # ------------------------------------------------------------------ #
-
- @patch("skyflow.vault.client.client.generate_bearer_token", return_value=("new_sa_token", None))
- @patch("skyflow.vault.client.client.is_expired", return_value=True)
- @patch("skyflow.vault.client.client.get_credentials")
- @patch("skyflow.vault.client.client.get_vault_url")
- @patch("skyflow.vault.client.client.VaultClient.initialize_api_client")
- def test_initialize_client_configuration_expired_token_no_reinit(
- self, mock_init_api_client, mock_get_vault_url, mock_get_credentials,
- mock_is_expired, mock_generate_bearer_token
- ):
- """Expired service account token is regenerated in-place; httpx client is NOT recreated."""
- mock_get_credentials.return_value = CREDENTIALS_WITH_PATH
- mock_get_vault_url.return_value = "https://test-vault-url.com"
-
- # Client already initialized — simulate warm state with an expired token
- self.vault_client._VaultClient__api_client = MagicMock()
- self.vault_client._VaultClient__is_static_token = False
- self.vault_client._VaultClient__bearer_token = "expired_sa_token"
- self.vault_client._VaultClient__credentials = CREDENTIALS_WITH_PATH
-
- self.vault_client.initialize_client_configuration()
-
- # Token was regenerated
- mock_generate_bearer_token.assert_called_once()
- self.assertEqual(
- self.vault_client._VaultClient__bearer_token, "new_sa_token"
- )
- # httpx client was NOT recreated
- mock_init_api_client.assert_not_called()
-
- # ------------------------------------------------------------------ #
- # initialize_client_configuration — config update forces reinit #
- # ------------------------------------------------------------------ #
-
- @patch("skyflow.vault.client.client.get_credentials")
- @patch("skyflow.vault.client.client.get_vault_url")
- @patch("skyflow.vault.client.client.VaultClient.initialize_api_client")
- def test_initialize_client_configuration_reinit_after_update_config(
- self, mock_init_api_client, mock_get_vault_url, mock_get_credentials
- ):
- """update_config() marks the client stale; next call must recreate it."""
- mock_get_credentials.return_value = CREDENTIALS_WITH_API_KEY
- mock_get_vault_url.return_value = "https://test-vault-url.com"
-
- # Simulate already-initialized client
- self.vault_client._VaultClient__api_client = MagicMock()
- self.vault_client._VaultClient__is_static_token = True
-
- self.vault_client.update_config({"cluster_id": "new_cluster"})
- self.vault_client.initialize_client_configuration()
-
- mock_get_credentials.assert_called_once()
- mock_get_vault_url.assert_called_once()
- mock_init_api_client.assert_called_once()
-
- # ------------------------------------------------------------------ #
- # initialize_api_client — lambda token provider #
- # ------------------------------------------------------------------ #
-
- @patch("skyflow.vault.client.client.Skyflow")
- def test_initialize_api_client_passes_callable_token(self, mock_skyflow):
- """initialize_api_client must pass a callable (lambda) as token, not a string."""
- self.vault_client.initialize_api_client("https://test-vault-url.com", "initial_token")
-
- args, kwargs = mock_skyflow.call_args
- self.assertEqual(kwargs["base_url"], "https://test-vault-url.com")
- self.assertTrue(callable(kwargs["token"]), "token must be a callable (lambda)")
-
- @patch("skyflow.vault.client.client.Skyflow")
- def test_initialize_api_client_lambda_returns_cached_bearer_token(self, mock_skyflow):
- """Lambda returns __bearer_token when it is set (interceptor behaviour)."""
- self.vault_client._VaultClient__bearer_token = "refreshed_token"
- self.vault_client.initialize_api_client("https://test-vault-url.com", "initial_token")
-
- _, kwargs = mock_skyflow.call_args
- self.assertEqual(kwargs["token"](), "refreshed_token")
-
- @patch("skyflow.vault.client.client.Skyflow")
- def test_initialize_api_client_lambda_falls_back_to_initial_token(self, mock_skyflow):
- """Lambda falls back to the initial token when __bearer_token is None."""
- self.vault_client._VaultClient__bearer_token = None
- self.vault_client.initialize_api_client("https://test-vault-url.com", "initial_token")
-
- _, kwargs = mock_skyflow.call_args
- self.assertEqual(kwargs["token"](), "initial_token")
-
- # ------------------------------------------------------------------ #
- # get_bearer_token #
- # ------------------------------------------------------------------ #
-
- def test_get_bearer_token_with_api_key(self):
- result = self.vault_client.get_bearer_token(CREDENTIALS_WITH_API_KEY)
- self.assertEqual(result, "dummy_api_key")
-
- def test_get_bearer_token_with_static_token(self):
- result = self.vault_client.get_bearer_token(CREDENTIALS_WITH_TOKEN)
- self.assertEqual(result, "dummy_static_token")
-
- @patch("skyflow.vault.client.client.generate_bearer_token", return_value=("sa_token", None))
- def test_get_bearer_token_generates_from_path_on_first_call(self, mock_generate):
- result = self.vault_client.get_bearer_token(CREDENTIALS_WITH_PATH)
- mock_generate.assert_called_once()
- self.assertEqual(result, "sa_token")
- self.assertEqual(self.vault_client._VaultClient__bearer_token, "sa_token")
-
- @patch("skyflow.vault.client.client.generate_bearer_token_from_creds", return_value=("sa_token_str", None))
- @patch("skyflow.vault.client.client.log_info")
- def test_get_bearer_token_generates_from_credentials_string(self, mock_log, mock_generate):
- result = self.vault_client.get_bearer_token(CREDENTIALS_WITH_STRING)
- mock_generate.assert_called_once()
- self.assertEqual(result, "sa_token_str")
-
- @patch("skyflow.vault.client.client.generate_bearer_token", return_value=("new_token", None))
- @patch("skyflow.vault.client.client.is_expired", return_value=True)
- @patch("skyflow.vault.client.client.log_info")
- def test_get_bearer_token_regenerates_on_expiry(self, mock_log, mock_is_expired, mock_generate):
- """Expired token is regenerated silently — no exception raised."""
- self.vault_client._VaultClient__bearer_token = "expired_token"
- result = self.vault_client.get_bearer_token(CREDENTIALS_WITH_PATH)
- mock_generate.assert_called_once()
- self.assertEqual(result, "new_token")
-
- @patch("skyflow.vault.client.client.generate_bearer_token")
- @patch("skyflow.vault.client.client.is_expired", return_value=False)
- @patch("skyflow.vault.client.client.log_info")
- def test_get_bearer_token_reuses_valid_cached_token(self, mock_log, mock_is_expired, mock_generate):
- """Valid cached token is reused without calling generate_bearer_token."""
- self.vault_client._VaultClient__bearer_token = "valid_token"
- result = self.vault_client.get_bearer_token(CREDENTIALS_WITH_PATH)
- mock_generate.assert_not_called()
- self.assertEqual(result, "valid_token")
-
- # ------------------------------------------------------------------ #
- # update_config #
- # ------------------------------------------------------------------ #
-
- def test_update_config_sets_flag(self):
- self.vault_client.update_config({"credentials": "new_credentials"})
- self.assertTrue(self.vault_client._VaultClient__is_config_updated)
- self.assertEqual(self.vault_client.get_config()["credentials"], "new_credentials")
-
- # ------------------------------------------------------------------ #
- # API accessor stubs #
- # ------------------------------------------------------------------ #
-
- def test_get_records_api(self):
- self.vault_client._VaultClient__api_client = MagicMock()
- self.assertIsNotNone(self.vault_client.get_records_api())
-
- def test_get_tokens_api(self):
- self.vault_client._VaultClient__api_client = MagicMock()
- self.assertIsNotNone(self.vault_client.get_tokens_api())
-
- def test_get_query_api(self):
- self.vault_client._VaultClient__api_client = MagicMock()
- self.assertIsNotNone(self.vault_client.get_query_api())
-
-
-if __name__ == "__main__":
- unittest.main()
\ No newline at end of file
diff --git a/requirements.txt b/v2/requirements.txt
similarity index 100%
rename from requirements.txt
rename to v2/requirements.txt
diff --git a/v2/setup.py b/v2/setup.py
new file mode 100644
index 00000000..4e7f78bc
--- /dev/null
+++ b/v2/setup.py
@@ -0,0 +1,85 @@
+'''
+ Copyright (c) 2022 Skyflow, Inc.
+'''
+import os
+import shutil
+import sys
+
+from setuptools import setup, find_packages
+from setuptools.command.build_py import build_py as _build_py
+
+
+if sys.version_info < (3, 9):
+ raise RuntimeError("skyflow requires Python 3.9+")
+current_version = '2.1.2'
+
+HERE = os.path.abspath(os.path.dirname(__file__))
+REPO_ROOT = os.path.dirname(HERE)
+COMMON_SRC = os.path.join(REPO_ROOT, 'common')
+
+with open(os.path.join(REPO_ROOT, 'README.md'), 'r', encoding='utf-8') as f:
+ long_description = f.read()
+
+# Anything under common/ that must never ride along into a built wheel.
+_COMMON_EXCLUDE_DIRS = {'__pycache__', '.pytest_cache', 'tests', '.mypy_cache'}
+_COMMON_EXCLUDE_FILES = {'setup.py', 'pyproject.toml', 'requirements.txt', '.gitignore'}
+_COMMON_EXCLUDE_SUFFIXES = ('.egg-info',)
+
+
+def _ignore_common_files(_directory, names):
+ ignored = set()
+ for name in names:
+ if name in _COMMON_EXCLUDE_DIRS or name in _COMMON_EXCLUDE_FILES:
+ ignored.add(name)
+ elif name.endswith(_COMMON_EXCLUDE_SUFFIXES):
+ ignored.add(name)
+ return ignored
+
+
+class CustomBuildPy(_build_py):
+ """SK-2938 Option C: bundles the sibling common/ source tree into this variant's wheel
+ (bdist_wheel/build --wheel only -- sdist doesn't invoke this and isn't supported here)."""
+
+ def run(self):
+ super().run()
+ dest = os.path.join(self.build_lib, 'common')
+ if os.path.exists(dest):
+ shutil.rmtree(dest)
+ shutil.copytree(COMMON_SRC, dest, ignore=_ignore_common_files)
+
+
+setup(
+ name='skyflow',
+ version=current_version,
+ author='Skyflow',
+ author_email='service-ops@skyflow.com',
+ packages=find_packages(where='.', exclude=['test*', 'samples*']),
+ # Ship PEP 561 markers so type checkers (mypy/pyright) see the SDK's types.
+ package_data={
+ 'skyflow': ['py.typed'],
+ 'skyflow.generated.rest': ['py.typed'],
+ },
+ cmdclass={'build_py': CustomBuildPy},
+ url='https://github.com/skyflowapi/skyflow-python/',
+ license='LICENSE',
+ description='Skyflow SDK for the Python programming language',
+ long_description=long_description,
+ long_description_content_type='text/markdown',
+ install_requires=[
+ 'pydantic >= 2.0.0',
+ 'typing-extensions >= 4.0.0',
+ 'PyJWT >= 2.12, < 3',
+ 'requests >= 2.28.0',
+ 'cryptography >= 44.0.2',
+ 'httpx >= 0.21.2',
+ 'python-dotenv >= 1.1.0, < 2',
+ ],
+ extras_require={
+ 'dev': [
+ 'codespell >= 2.4.1',
+ 'ruff >= 0.9.0',
+ 'pre-commit >= 4.3.0',
+ ]
+ },
+ python_requires=">=3.9",
+)
diff --git a/v2/skyflow/__init__.py b/v2/skyflow/__init__.py
new file mode 100644
index 00000000..fc02764f
--- /dev/null
+++ b/v2/skyflow/__init__.py
@@ -0,0 +1,2 @@
+from .utils import LogLevel, Env
+from .client import Skyflow
diff --git a/v2/skyflow/client/__init__.py b/v2/skyflow/client/__init__.py
new file mode 100644
index 00000000..246ca2f6
--- /dev/null
+++ b/v2/skyflow/client/__init__.py
@@ -0,0 +1 @@
+from .skyflow import Skyflow
diff --git a/v2/skyflow/client/skyflow.py b/v2/skyflow/client/skyflow.py
new file mode 100644
index 00000000..c464b2f7
--- /dev/null
+++ b/v2/skyflow/client/skyflow.py
@@ -0,0 +1,17 @@
+from common.client.utils import make_skyflow_class
+from skyflow.utils import SkyflowMessages
+from skyflow.utils.logger import set_active_log_level
+from skyflow.utils.validations import validate_connection_config, validate_update_connection_config
+from skyflow.vault.client.client import VaultClient
+from skyflow.vault.controller import VaultController, Connection, Detect
+
+Skyflow = make_skyflow_class(
+ vault_client_cls=VaultClient,
+ vault_controller_cls=VaultController,
+ connection_cls=Connection,
+ detect_cls=Detect,
+ skyflow_messages=SkyflowMessages,
+ validate_connection_config=validate_connection_config,
+ validate_update_connection_config=validate_update_connection_config,
+ set_active_log_level=set_active_log_level,
+)
diff --git a/v2/skyflow/error/__init__.py b/v2/skyflow/error/__init__.py
new file mode 100644
index 00000000..a02e980c
--- /dev/null
+++ b/v2/skyflow/error/__init__.py
@@ -0,0 +1,3 @@
+from common.errors import SkyflowError
+
+__all__ = ["SkyflowError"]
\ No newline at end of file
diff --git a/v2/skyflow/generated/__init__.py b/v2/skyflow/generated/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/skyflow/generated/rest/__init__.py b/v2/skyflow/generated/rest/__init__.py
similarity index 100%
rename from skyflow/generated/rest/__init__.py
rename to v2/skyflow/generated/rest/__init__.py
diff --git a/skyflow/generated/rest/audit/__init__.py b/v2/skyflow/generated/rest/audit/__init__.py
similarity index 100%
rename from skyflow/generated/rest/audit/__init__.py
rename to v2/skyflow/generated/rest/audit/__init__.py
diff --git a/skyflow/generated/rest/audit/client.py b/v2/skyflow/generated/rest/audit/client.py
similarity index 100%
rename from skyflow/generated/rest/audit/client.py
rename to v2/skyflow/generated/rest/audit/client.py
diff --git a/skyflow/generated/rest/audit/raw_client.py b/v2/skyflow/generated/rest/audit/raw_client.py
similarity index 100%
rename from skyflow/generated/rest/audit/raw_client.py
rename to v2/skyflow/generated/rest/audit/raw_client.py
diff --git a/skyflow/generated/rest/audit/types/__init__.py b/v2/skyflow/generated/rest/audit/types/__init__.py
similarity index 100%
rename from skyflow/generated/rest/audit/types/__init__.py
rename to v2/skyflow/generated/rest/audit/types/__init__.py
diff --git a/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_action_type.py b/v2/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_action_type.py
similarity index 100%
rename from skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_action_type.py
rename to v2/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_action_type.py
diff --git a/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_context_access_type.py b/v2/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_context_access_type.py
similarity index 100%
rename from skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_context_access_type.py
rename to v2/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_context_access_type.py
diff --git a/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_context_actor_type.py b/v2/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_context_actor_type.py
similarity index 100%
rename from skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_context_actor_type.py
rename to v2/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_context_actor_type.py
diff --git a/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_context_auth_mode.py b/v2/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_context_auth_mode.py
similarity index 100%
rename from skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_context_auth_mode.py
rename to v2/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_context_auth_mode.py
diff --git a/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_resource_type.py b/v2/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_resource_type.py
similarity index 100%
rename from skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_resource_type.py
rename to v2/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_filter_ops_resource_type.py
diff --git a/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_sort_ops_order_by.py b/v2/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_sort_ops_order_by.py
similarity index 100%
rename from skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_sort_ops_order_by.py
rename to v2/skyflow/generated/rest/audit/types/audit_service_list_audit_events_request_sort_ops_order_by.py
diff --git a/skyflow/generated/rest/query/__init__.py b/v2/skyflow/generated/rest/authentication/__init__.py
similarity index 100%
rename from skyflow/generated/rest/query/__init__.py
rename to v2/skyflow/generated/rest/authentication/__init__.py
diff --git a/skyflow/generated/rest/authentication/client.py b/v2/skyflow/generated/rest/authentication/client.py
similarity index 100%
rename from skyflow/generated/rest/authentication/client.py
rename to v2/skyflow/generated/rest/authentication/client.py
diff --git a/skyflow/generated/rest/authentication/raw_client.py b/v2/skyflow/generated/rest/authentication/raw_client.py
similarity index 100%
rename from skyflow/generated/rest/authentication/raw_client.py
rename to v2/skyflow/generated/rest/authentication/raw_client.py
diff --git a/skyflow/generated/rest/tokens/__init__.py b/v2/skyflow/generated/rest/bin_lookup/__init__.py
similarity index 100%
rename from skyflow/generated/rest/tokens/__init__.py
rename to v2/skyflow/generated/rest/bin_lookup/__init__.py
diff --git a/skyflow/generated/rest/bin_lookup/client.py b/v2/skyflow/generated/rest/bin_lookup/client.py
similarity index 100%
rename from skyflow/generated/rest/bin_lookup/client.py
rename to v2/skyflow/generated/rest/bin_lookup/client.py
diff --git a/skyflow/generated/rest/bin_lookup/raw_client.py b/v2/skyflow/generated/rest/bin_lookup/raw_client.py
similarity index 100%
rename from skyflow/generated/rest/bin_lookup/raw_client.py
rename to v2/skyflow/generated/rest/bin_lookup/raw_client.py
diff --git a/skyflow/generated/rest/client.py b/v2/skyflow/generated/rest/client.py
similarity index 100%
rename from skyflow/generated/rest/client.py
rename to v2/skyflow/generated/rest/client.py
diff --git a/v2/skyflow/generated/rest/core/__init__.py b/v2/skyflow/generated/rest/core/__init__.py
new file mode 100644
index 00000000..31bbb818
--- /dev/null
+++ b/v2/skyflow/generated/rest/core/__init__.py
@@ -0,0 +1,52 @@
+# This file was auto-generated by Fern from our API Definition.
+
+# isort: skip_file
+
+from .api_error import ApiError
+from .client_wrapper import AsyncClientWrapper, BaseClientWrapper, SyncClientWrapper
+from .datetime_utils import serialize_datetime
+from .file import File, convert_file_dict_to_httpx_tuples, with_content_type
+from .http_client import AsyncHttpClient, HttpClient
+from .http_response import AsyncHttpResponse, HttpResponse
+from .jsonable_encoder import jsonable_encoder
+from .pydantic_utilities import (
+ IS_PYDANTIC_V2,
+ UniversalBaseModel,
+ UniversalRootModel,
+ parse_obj_as,
+ universal_field_validator,
+ universal_root_validator,
+ update_forward_refs,
+)
+from .query_encoder import encode_query
+from .remove_none_from_dict import remove_none_from_dict
+from .request_options import RequestOptions
+from .serialization import FieldMetadata, convert_and_respect_annotation_metadata
+
+__all__ = [
+ "ApiError",
+ "AsyncClientWrapper",
+ "AsyncHttpClient",
+ "AsyncHttpResponse",
+ "BaseClientWrapper",
+ "FieldMetadata",
+ "File",
+ "HttpClient",
+ "HttpResponse",
+ "IS_PYDANTIC_V2",
+ "RequestOptions",
+ "SyncClientWrapper",
+ "UniversalBaseModel",
+ "UniversalRootModel",
+ "convert_and_respect_annotation_metadata",
+ "convert_file_dict_to_httpx_tuples",
+ "encode_query",
+ "jsonable_encoder",
+ "parse_obj_as",
+ "remove_none_from_dict",
+ "serialize_datetime",
+ "universal_field_validator",
+ "universal_root_validator",
+ "update_forward_refs",
+ "with_content_type",
+]
diff --git a/v2/skyflow/generated/rest/core/api_error.py b/v2/skyflow/generated/rest/core/api_error.py
new file mode 100644
index 00000000..6f850a60
--- /dev/null
+++ b/v2/skyflow/generated/rest/core/api_error.py
@@ -0,0 +1,23 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from typing import Any, Dict, Optional
+
+
+class ApiError(Exception):
+ headers: Optional[Dict[str, str]]
+ status_code: Optional[int]
+ body: Any
+
+ def __init__(
+ self,
+ *,
+ headers: Optional[Dict[str, str]] = None,
+ status_code: Optional[int] = None,
+ body: Any = None,
+ ) -> None:
+ self.headers = headers
+ self.status_code = status_code
+ self.body = body
+
+ def __str__(self) -> str:
+ return f"headers: {self.headers}, status_code: {self.status_code}, body: {self.body}"
diff --git a/skyflow/generated/rest/core/client_wrapper.py b/v2/skyflow/generated/rest/core/client_wrapper.py
similarity index 100%
rename from skyflow/generated/rest/core/client_wrapper.py
rename to v2/skyflow/generated/rest/core/client_wrapper.py
diff --git a/v2/skyflow/generated/rest/core/datetime_utils.py b/v2/skyflow/generated/rest/core/datetime_utils.py
new file mode 100644
index 00000000..7c9864a9
--- /dev/null
+++ b/v2/skyflow/generated/rest/core/datetime_utils.py
@@ -0,0 +1,28 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import datetime as dt
+
+
+def serialize_datetime(v: dt.datetime) -> str:
+ """
+ Serialize a datetime including timezone info.
+
+ Uses the timezone info provided if present, otherwise uses the current runtime's timezone info.
+
+ UTC datetimes end in "Z" while all other timezones are represented as offset from UTC, e.g. +05:00.
+ """
+
+ def _serialize_zoned_datetime(v: dt.datetime) -> str:
+ if v.tzinfo is not None and v.tzinfo.tzname(None) == dt.timezone.utc.tzname(None):
+ # UTC is a special case where we use "Z" at the end instead of "+00:00"
+ return v.isoformat().replace("+00:00", "Z")
+ else:
+ # Delegate to the typical +/- offset format
+ return v.isoformat()
+
+ if v.tzinfo is not None:
+ return _serialize_zoned_datetime(v)
+ else:
+ local_tz = dt.datetime.now().astimezone().tzinfo
+ localized_dt = v.replace(tzinfo=local_tz)
+ return _serialize_zoned_datetime(localized_dt)
diff --git a/v2/skyflow/generated/rest/core/file.py b/v2/skyflow/generated/rest/core/file.py
new file mode 100644
index 00000000..44b0d27c
--- /dev/null
+++ b/v2/skyflow/generated/rest/core/file.py
@@ -0,0 +1,67 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from typing import IO, Dict, List, Mapping, Optional, Tuple, Union, cast
+
+# File typing inspired by the flexibility of types within the httpx library
+# https://github.com/encode/httpx/blob/master/httpx/_types.py
+FileContent = Union[IO[bytes], bytes, str]
+File = Union[
+ # file (or bytes)
+ FileContent,
+ # (filename, file (or bytes))
+ Tuple[Optional[str], FileContent],
+ # (filename, file (or bytes), content_type)
+ Tuple[Optional[str], FileContent, Optional[str]],
+ # (filename, file (or bytes), content_type, headers)
+ Tuple[
+ Optional[str],
+ FileContent,
+ Optional[str],
+ Mapping[str, str],
+ ],
+]
+
+
+def convert_file_dict_to_httpx_tuples(
+ d: Dict[str, Union[File, List[File]]],
+) -> List[Tuple[str, File]]:
+ """
+ The format we use is a list of tuples, where the first element is the
+ name of the file and the second is the file object. Typically HTTPX wants
+ a dict, but to be able to send lists of files, you have to use the list
+ approach (which also works for non-lists)
+ https://github.com/encode/httpx/pull/1032
+ """
+
+ httpx_tuples = []
+ for key, file_like in d.items():
+ if isinstance(file_like, list):
+ for file_like_item in file_like:
+ httpx_tuples.append((key, file_like_item))
+ else:
+ httpx_tuples.append((key, file_like))
+ return httpx_tuples
+
+
+def with_content_type(*, file: File, default_content_type: str) -> File:
+ """
+ This function resolves to the file's content type, if provided, and defaults
+ to the default_content_type value if not.
+ """
+ if isinstance(file, tuple):
+ if len(file) == 2:
+ filename, content = cast(Tuple[Optional[str], FileContent], file) # type: ignore
+ return (filename, content, default_content_type)
+ elif len(file) == 3:
+ filename, content, file_content_type = cast(Tuple[Optional[str], FileContent, Optional[str]], file) # type: ignore
+ out_content_type = file_content_type or default_content_type
+ return (filename, content, out_content_type)
+ elif len(file) == 4:
+ filename, content, file_content_type, headers = cast( # type: ignore
+ Tuple[Optional[str], FileContent, Optional[str], Mapping[str, str]], file
+ )
+ out_content_type = file_content_type or default_content_type
+ return (filename, content, out_content_type, headers)
+ else:
+ raise ValueError(f"Unexpected tuple length: {len(file)}")
+ return (None, file, default_content_type)
diff --git a/v2/skyflow/generated/rest/core/force_multipart.py b/v2/skyflow/generated/rest/core/force_multipart.py
new file mode 100644
index 00000000..ae24ccff
--- /dev/null
+++ b/v2/skyflow/generated/rest/core/force_multipart.py
@@ -0,0 +1,16 @@
+# This file was auto-generated by Fern from our API Definition.
+
+
+class ForceMultipartDict(dict):
+ """
+ A dictionary subclass that always evaluates to True in boolean contexts.
+
+ This is used to force multipart/form-data encoding in HTTP requests even when
+ the dictionary is empty, which would normally evaluate to False.
+ """
+
+ def __bool__(self):
+ return True
+
+
+FORCE_MULTIPART = ForceMultipartDict()
diff --git a/v2/skyflow/generated/rest/core/http_client.py b/v2/skyflow/generated/rest/core/http_client.py
new file mode 100644
index 00000000..e4173f99
--- /dev/null
+++ b/v2/skyflow/generated/rest/core/http_client.py
@@ -0,0 +1,543 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import asyncio
+import email.utils
+import re
+import time
+import typing
+import urllib.parse
+from contextlib import asynccontextmanager, contextmanager
+from random import random
+
+import httpx
+from .file import File, convert_file_dict_to_httpx_tuples
+from .force_multipart import FORCE_MULTIPART
+from .jsonable_encoder import jsonable_encoder
+from .query_encoder import encode_query
+from .remove_none_from_dict import remove_none_from_dict
+from .request_options import RequestOptions
+from httpx._types import RequestFiles
+
+INITIAL_RETRY_DELAY_SECONDS = 0.5
+MAX_RETRY_DELAY_SECONDS = 10
+MAX_RETRY_DELAY_SECONDS_FROM_HEADER = 30
+
+
+def _parse_retry_after(response_headers: httpx.Headers) -> typing.Optional[float]:
+ """
+ This function parses the `Retry-After` header in a HTTP response and returns the number of seconds to wait.
+
+ Inspired by the urllib3 retry implementation.
+ """
+ retry_after_ms = response_headers.get("retry-after-ms")
+ if retry_after_ms is not None:
+ try:
+ return int(retry_after_ms) / 1000 if retry_after_ms > 0 else 0
+ except Exception:
+ pass
+
+ retry_after = response_headers.get("retry-after")
+ if retry_after is None:
+ return None
+
+ # Attempt to parse the header as an int.
+ if re.match(r"^\s*[0-9]+\s*$", retry_after):
+ seconds = float(retry_after)
+ # Fallback to parsing it as a date.
+ else:
+ retry_date_tuple = email.utils.parsedate_tz(retry_after)
+ if retry_date_tuple is None:
+ return None
+ if retry_date_tuple[9] is None: # Python 2
+ # Assume UTC if no timezone was specified
+ # On Python2.7, parsedate_tz returns None for a timezone offset
+ # instead of 0 if no timezone is given, where mktime_tz treats
+ # a None timezone offset as local time.
+ retry_date_tuple = retry_date_tuple[:9] + (0,) + retry_date_tuple[10:]
+
+ retry_date = email.utils.mktime_tz(retry_date_tuple)
+ seconds = retry_date - time.time()
+
+ if seconds < 0:
+ seconds = 0
+
+ return seconds
+
+
+def _retry_timeout(response: httpx.Response, retries: int) -> float:
+ """
+ Determine the amount of time to wait before retrying a request.
+ This function begins by trying to parse a retry-after header from the response, and then proceeds to use exponential backoff
+ with a jitter to determine the number of seconds to wait.
+ """
+
+ # If the API asks us to wait a certain amount of time (and it's a reasonable amount), just do what it says.
+ retry_after = _parse_retry_after(response.headers)
+ if retry_after is not None and retry_after <= MAX_RETRY_DELAY_SECONDS_FROM_HEADER:
+ return retry_after
+
+ # Apply exponential backoff, capped at MAX_RETRY_DELAY_SECONDS.
+ retry_delay = min(INITIAL_RETRY_DELAY_SECONDS * pow(2.0, retries), MAX_RETRY_DELAY_SECONDS)
+
+ # Add a randomness / jitter to the retry delay to avoid overwhelming the server with retries.
+ timeout = retry_delay * (1 - 0.25 * random())
+ return timeout if timeout >= 0 else 0
+
+
+def _should_retry(response: httpx.Response) -> bool:
+ retryable_400s = [429, 408, 409]
+ return response.status_code >= 500 or response.status_code in retryable_400s
+
+
+def remove_omit_from_dict(
+ original: typing.Dict[str, typing.Optional[typing.Any]],
+ omit: typing.Optional[typing.Any],
+) -> typing.Dict[str, typing.Any]:
+ if omit is None:
+ return original
+ new: typing.Dict[str, typing.Any] = {}
+ for key, value in original.items():
+ if value is not omit:
+ new[key] = value
+ return new
+
+
+def maybe_filter_request_body(
+ data: typing.Optional[typing.Any],
+ request_options: typing.Optional[RequestOptions],
+ omit: typing.Optional[typing.Any],
+) -> typing.Optional[typing.Any]:
+ if data is None:
+ return (
+ jsonable_encoder(request_options.get("additional_body_parameters", {})) or {}
+ if request_options is not None
+ else None
+ )
+ elif not isinstance(data, typing.Mapping):
+ data_content = jsonable_encoder(data)
+ else:
+ data_content = {
+ **(jsonable_encoder(remove_omit_from_dict(data, omit))), # type: ignore
+ **(
+ jsonable_encoder(request_options.get("additional_body_parameters", {})) or {}
+ if request_options is not None
+ else {}
+ ),
+ }
+ return data_content
+
+
+# Abstracted out for testing purposes
+def get_request_body(
+ *,
+ json: typing.Optional[typing.Any],
+ data: typing.Optional[typing.Any],
+ request_options: typing.Optional[RequestOptions],
+ omit: typing.Optional[typing.Any],
+) -> typing.Tuple[typing.Optional[typing.Any], typing.Optional[typing.Any]]:
+ json_body = None
+ data_body = None
+ if data is not None:
+ data_body = maybe_filter_request_body(data, request_options, omit)
+ else:
+ # If both data and json are None, we send json data in the event extra properties are specified
+ json_body = maybe_filter_request_body(json, request_options, omit)
+
+ # If you have an empty JSON body, you should just send None
+ return (json_body if json_body != {} else None), data_body if data_body != {} else None
+
+
+class HttpClient:
+ def __init__(
+ self,
+ *,
+ httpx_client: httpx.Client,
+ base_timeout: typing.Callable[[], typing.Optional[float]],
+ base_headers: typing.Callable[[], typing.Dict[str, str]],
+ base_url: typing.Optional[typing.Callable[[], str]] = None,
+ ):
+ self.base_url = base_url
+ self.base_timeout = base_timeout
+ self.base_headers = base_headers
+ self.httpx_client = httpx_client
+
+ def get_base_url(self, maybe_base_url: typing.Optional[str]) -> str:
+ base_url = maybe_base_url
+ if self.base_url is not None and base_url is None:
+ base_url = self.base_url()
+
+ if base_url is None:
+ raise ValueError("A base_url is required to make this request, please provide one and try again.")
+ return base_url
+
+ def request(
+ self,
+ path: typing.Optional[str] = None,
+ *,
+ method: str,
+ base_url: typing.Optional[str] = None,
+ params: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ json: typing.Optional[typing.Any] = None,
+ data: typing.Optional[typing.Any] = None,
+ content: typing.Optional[typing.Union[bytes, typing.Iterator[bytes], typing.AsyncIterator[bytes]]] = None,
+ files: typing.Optional[
+ typing.Union[
+ typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]],
+ typing.List[typing.Tuple[str, File]],
+ ]
+ ] = None,
+ headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ request_options: typing.Optional[RequestOptions] = None,
+ retries: int = 2,
+ omit: typing.Optional[typing.Any] = None,
+ force_multipart: typing.Optional[bool] = None,
+ ) -> httpx.Response:
+ base_url = self.get_base_url(base_url)
+ timeout = (
+ request_options.get("timeout_in_seconds")
+ if request_options is not None and request_options.get("timeout_in_seconds") is not None
+ else self.base_timeout()
+ )
+
+ json_body, data_body = get_request_body(json=json, data=data, request_options=request_options, omit=omit)
+
+ request_files: typing.Optional[RequestFiles] = (
+ convert_file_dict_to_httpx_tuples(remove_omit_from_dict(remove_none_from_dict(files), omit))
+ if (files is not None and files is not omit and isinstance(files, dict))
+ else None
+ )
+
+ if (request_files is None or len(request_files) == 0) and force_multipart:
+ request_files = FORCE_MULTIPART
+
+ response = self.httpx_client.request(
+ method=method,
+ url=urllib.parse.urljoin(f"{base_url}/", path),
+ headers=jsonable_encoder(
+ remove_none_from_dict(
+ {
+ **self.base_headers(),
+ **(headers if headers is not None else {}),
+ **(request_options.get("additional_headers", {}) or {} if request_options is not None else {}),
+ }
+ )
+ ),
+ params=encode_query(
+ jsonable_encoder(
+ remove_none_from_dict(
+ remove_omit_from_dict(
+ {
+ **(params if params is not None else {}),
+ **(
+ request_options.get("additional_query_parameters", {}) or {}
+ if request_options is not None
+ else {}
+ ),
+ },
+ omit,
+ )
+ )
+ )
+ ),
+ json=json_body,
+ data=data_body,
+ content=content,
+ files=request_files,
+ timeout=timeout,
+ )
+
+ max_retries: int = request_options.get("max_retries", 0) if request_options is not None else 0
+ if _should_retry(response=response):
+ if max_retries > retries:
+ time.sleep(_retry_timeout(response=response, retries=retries))
+ return self.request(
+ path=path,
+ method=method,
+ base_url=base_url,
+ params=params,
+ json=json,
+ content=content,
+ files=files,
+ headers=headers,
+ request_options=request_options,
+ retries=retries + 1,
+ omit=omit,
+ )
+
+ return response
+
+ @contextmanager
+ def stream(
+ self,
+ path: typing.Optional[str] = None,
+ *,
+ method: str,
+ base_url: typing.Optional[str] = None,
+ params: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ json: typing.Optional[typing.Any] = None,
+ data: typing.Optional[typing.Any] = None,
+ content: typing.Optional[typing.Union[bytes, typing.Iterator[bytes], typing.AsyncIterator[bytes]]] = None,
+ files: typing.Optional[
+ typing.Union[
+ typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]],
+ typing.List[typing.Tuple[str, File]],
+ ]
+ ] = None,
+ headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ request_options: typing.Optional[RequestOptions] = None,
+ retries: int = 2,
+ omit: typing.Optional[typing.Any] = None,
+ force_multipart: typing.Optional[bool] = None,
+ ) -> typing.Iterator[httpx.Response]:
+ base_url = self.get_base_url(base_url)
+ timeout = (
+ request_options.get("timeout_in_seconds")
+ if request_options is not None and request_options.get("timeout_in_seconds") is not None
+ else self.base_timeout()
+ )
+
+ request_files: typing.Optional[RequestFiles] = (
+ convert_file_dict_to_httpx_tuples(remove_omit_from_dict(remove_none_from_dict(files), omit))
+ if (files is not None and files is not omit and isinstance(files, dict))
+ else None
+ )
+
+ if (request_files is None or len(request_files) == 0) and force_multipart:
+ request_files = FORCE_MULTIPART
+
+ json_body, data_body = get_request_body(json=json, data=data, request_options=request_options, omit=omit)
+
+ with self.httpx_client.stream(
+ method=method,
+ url=urllib.parse.urljoin(f"{base_url}/", path),
+ headers=jsonable_encoder(
+ remove_none_from_dict(
+ {
+ **self.base_headers(),
+ **(headers if headers is not None else {}),
+ **(request_options.get("additional_headers", {}) if request_options is not None else {}),
+ }
+ )
+ ),
+ params=encode_query(
+ jsonable_encoder(
+ remove_none_from_dict(
+ remove_omit_from_dict(
+ {
+ **(params if params is not None else {}),
+ **(
+ request_options.get("additional_query_parameters", {})
+ if request_options is not None
+ else {}
+ ),
+ },
+ omit,
+ )
+ )
+ )
+ ),
+ json=json_body,
+ data=data_body,
+ content=content,
+ files=request_files,
+ timeout=timeout,
+ ) as stream:
+ yield stream
+
+
+class AsyncHttpClient:
+ def __init__(
+ self,
+ *,
+ httpx_client: httpx.AsyncClient,
+ base_timeout: typing.Callable[[], typing.Optional[float]],
+ base_headers: typing.Callable[[], typing.Dict[str, str]],
+ base_url: typing.Optional[typing.Callable[[], str]] = None,
+ ):
+ self.base_url = base_url
+ self.base_timeout = base_timeout
+ self.base_headers = base_headers
+ self.httpx_client = httpx_client
+
+ def get_base_url(self, maybe_base_url: typing.Optional[str]) -> str:
+ base_url = maybe_base_url
+ if self.base_url is not None and base_url is None:
+ base_url = self.base_url()
+
+ if base_url is None:
+ raise ValueError("A base_url is required to make this request, please provide one and try again.")
+ return base_url
+
+ async def request(
+ self,
+ path: typing.Optional[str] = None,
+ *,
+ method: str,
+ base_url: typing.Optional[str] = None,
+ params: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ json: typing.Optional[typing.Any] = None,
+ data: typing.Optional[typing.Any] = None,
+ content: typing.Optional[typing.Union[bytes, typing.Iterator[bytes], typing.AsyncIterator[bytes]]] = None,
+ files: typing.Optional[
+ typing.Union[
+ typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]],
+ typing.List[typing.Tuple[str, File]],
+ ]
+ ] = None,
+ headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ request_options: typing.Optional[RequestOptions] = None,
+ retries: int = 2,
+ omit: typing.Optional[typing.Any] = None,
+ force_multipart: typing.Optional[bool] = None,
+ ) -> httpx.Response:
+ base_url = self.get_base_url(base_url)
+ timeout = (
+ request_options.get("timeout_in_seconds")
+ if request_options is not None and request_options.get("timeout_in_seconds") is not None
+ else self.base_timeout()
+ )
+
+ request_files: typing.Optional[RequestFiles] = (
+ convert_file_dict_to_httpx_tuples(remove_omit_from_dict(remove_none_from_dict(files), omit))
+ if (files is not None and files is not omit and isinstance(files, dict))
+ else None
+ )
+
+ if (request_files is None or len(request_files) == 0) and force_multipart:
+ request_files = FORCE_MULTIPART
+
+ json_body, data_body = get_request_body(json=json, data=data, request_options=request_options, omit=omit)
+
+ # Add the input to each of these and do None-safety checks
+ response = await self.httpx_client.request(
+ method=method,
+ url=urllib.parse.urljoin(f"{base_url}/", path),
+ headers=jsonable_encoder(
+ remove_none_from_dict(
+ {
+ **self.base_headers(),
+ **(headers if headers is not None else {}),
+ **(request_options.get("additional_headers", {}) or {} if request_options is not None else {}),
+ }
+ )
+ ),
+ params=encode_query(
+ jsonable_encoder(
+ remove_none_from_dict(
+ remove_omit_from_dict(
+ {
+ **(params if params is not None else {}),
+ **(
+ request_options.get("additional_query_parameters", {}) or {}
+ if request_options is not None
+ else {}
+ ),
+ },
+ omit,
+ )
+ )
+ )
+ ),
+ json=json_body,
+ data=data_body,
+ content=content,
+ files=request_files,
+ timeout=timeout,
+ )
+
+ max_retries: int = request_options.get("max_retries", 0) if request_options is not None else 0
+ if _should_retry(response=response):
+ if max_retries > retries:
+ await asyncio.sleep(_retry_timeout(response=response, retries=retries))
+ return await self.request(
+ path=path,
+ method=method,
+ base_url=base_url,
+ params=params,
+ json=json,
+ content=content,
+ files=files,
+ headers=headers,
+ request_options=request_options,
+ retries=retries + 1,
+ omit=omit,
+ )
+ return response
+
+ @asynccontextmanager
+ async def stream(
+ self,
+ path: typing.Optional[str] = None,
+ *,
+ method: str,
+ base_url: typing.Optional[str] = None,
+ params: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ json: typing.Optional[typing.Any] = None,
+ data: typing.Optional[typing.Any] = None,
+ content: typing.Optional[typing.Union[bytes, typing.Iterator[bytes], typing.AsyncIterator[bytes]]] = None,
+ files: typing.Optional[
+ typing.Union[
+ typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]],
+ typing.List[typing.Tuple[str, File]],
+ ]
+ ] = None,
+ headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
+ request_options: typing.Optional[RequestOptions] = None,
+ retries: int = 2,
+ omit: typing.Optional[typing.Any] = None,
+ force_multipart: typing.Optional[bool] = None,
+ ) -> typing.AsyncIterator[httpx.Response]:
+ base_url = self.get_base_url(base_url)
+ timeout = (
+ request_options.get("timeout_in_seconds")
+ if request_options is not None and request_options.get("timeout_in_seconds") is not None
+ else self.base_timeout()
+ )
+
+ request_files: typing.Optional[RequestFiles] = (
+ convert_file_dict_to_httpx_tuples(remove_omit_from_dict(remove_none_from_dict(files), omit))
+ if (files is not None and files is not omit and isinstance(files, dict))
+ else None
+ )
+
+ if (request_files is None or len(request_files) == 0) and force_multipart:
+ request_files = FORCE_MULTIPART
+
+ json_body, data_body = get_request_body(json=json, data=data, request_options=request_options, omit=omit)
+
+ async with self.httpx_client.stream(
+ method=method,
+ url=urllib.parse.urljoin(f"{base_url}/", path),
+ headers=jsonable_encoder(
+ remove_none_from_dict(
+ {
+ **self.base_headers(),
+ **(headers if headers is not None else {}),
+ **(request_options.get("additional_headers", {}) if request_options is not None else {}),
+ }
+ )
+ ),
+ params=encode_query(
+ jsonable_encoder(
+ remove_none_from_dict(
+ remove_omit_from_dict(
+ {
+ **(params if params is not None else {}),
+ **(
+ request_options.get("additional_query_parameters", {})
+ if request_options is not None
+ else {}
+ ),
+ },
+ omit=omit,
+ )
+ )
+ )
+ ),
+ json=json_body,
+ data=data_body,
+ content=content,
+ files=request_files,
+ timeout=timeout,
+ ) as stream:
+ yield stream
diff --git a/v2/skyflow/generated/rest/core/http_response.py b/v2/skyflow/generated/rest/core/http_response.py
new file mode 100644
index 00000000..48a1798a
--- /dev/null
+++ b/v2/skyflow/generated/rest/core/http_response.py
@@ -0,0 +1,55 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from typing import Dict, Generic, TypeVar
+
+import httpx
+
+T = TypeVar("T")
+"""Generic to represent the underlying type of the data wrapped by the HTTP response."""
+
+
+class BaseHttpResponse:
+ """Minimalist HTTP response wrapper that exposes response headers."""
+
+ _response: httpx.Response
+
+ def __init__(self, response: httpx.Response):
+ self._response = response
+
+ @property
+ def headers(self) -> Dict[str, str]:
+ return dict(self._response.headers)
+
+
+class HttpResponse(Generic[T], BaseHttpResponse):
+ """HTTP response wrapper that exposes response headers and data."""
+
+ _data: T
+
+ def __init__(self, response: httpx.Response, data: T):
+ super().__init__(response)
+ self._data = data
+
+ @property
+ def data(self) -> T:
+ return self._data
+
+ def close(self) -> None:
+ self._response.close()
+
+
+class AsyncHttpResponse(Generic[T], BaseHttpResponse):
+ """HTTP response wrapper that exposes response headers and data."""
+
+ _data: T
+
+ def __init__(self, response: httpx.Response, data: T):
+ super().__init__(response)
+ self._data = data
+
+ @property
+ def data(self) -> T:
+ return self._data
+
+ async def close(self) -> None:
+ await self._response.aclose()
diff --git a/v2/skyflow/generated/rest/core/jsonable_encoder.py b/v2/skyflow/generated/rest/core/jsonable_encoder.py
new file mode 100644
index 00000000..afee3662
--- /dev/null
+++ b/v2/skyflow/generated/rest/core/jsonable_encoder.py
@@ -0,0 +1,100 @@
+# This file was auto-generated by Fern from our API Definition.
+
+"""
+jsonable_encoder converts a Python object to a JSON-friendly dict
+(e.g. datetimes to strings, Pydantic models to dicts).
+
+Taken from FastAPI, and made a bit simpler
+https://github.com/tiangolo/fastapi/blob/master/fastapi/encoders.py
+"""
+
+import base64
+import dataclasses
+import datetime as dt
+from enum import Enum
+from pathlib import PurePath
+from types import GeneratorType
+from typing import Any, Callable, Dict, List, Optional, Set, Union
+
+import pydantic
+from .datetime_utils import serialize_datetime
+from .pydantic_utilities import (
+ IS_PYDANTIC_V2,
+ encode_by_type,
+ to_jsonable_with_fallback,
+)
+
+SetIntStr = Set[Union[int, str]]
+DictIntStrAny = Dict[Union[int, str], Any]
+
+
+def jsonable_encoder(obj: Any, custom_encoder: Optional[Dict[Any, Callable[[Any], Any]]] = None) -> Any:
+ custom_encoder = custom_encoder or {}
+ if custom_encoder:
+ if type(obj) in custom_encoder:
+ return custom_encoder[type(obj)](obj)
+ else:
+ for encoder_type, encoder_instance in custom_encoder.items():
+ if isinstance(obj, encoder_type):
+ return encoder_instance(obj)
+ if isinstance(obj, pydantic.BaseModel):
+ if IS_PYDANTIC_V2:
+ encoder = getattr(obj.model_config, "json_encoders", {}) # type: ignore # Pydantic v2
+ else:
+ encoder = getattr(obj.__config__, "json_encoders", {}) # type: ignore # Pydantic v1
+ if custom_encoder:
+ encoder.update(custom_encoder)
+ obj_dict = obj.dict(by_alias=True)
+ if "__root__" in obj_dict:
+ obj_dict = obj_dict["__root__"]
+ if "root" in obj_dict:
+ obj_dict = obj_dict["root"]
+ return jsonable_encoder(obj_dict, custom_encoder=encoder)
+ if dataclasses.is_dataclass(obj):
+ obj_dict = dataclasses.asdict(obj) # type: ignore
+ return jsonable_encoder(obj_dict, custom_encoder=custom_encoder)
+ if isinstance(obj, bytes):
+ return base64.b64encode(obj).decode("utf-8")
+ if isinstance(obj, Enum):
+ return obj.value
+ if isinstance(obj, PurePath):
+ return str(obj)
+ if isinstance(obj, (str, int, float, type(None))):
+ return obj
+ if isinstance(obj, dt.datetime):
+ return serialize_datetime(obj)
+ if isinstance(obj, dt.date):
+ return str(obj)
+ if isinstance(obj, dict):
+ encoded_dict = {}
+ allowed_keys = set(obj.keys())
+ for key, value in obj.items():
+ if key in allowed_keys:
+ encoded_key = jsonable_encoder(key, custom_encoder=custom_encoder)
+ encoded_value = jsonable_encoder(value, custom_encoder=custom_encoder)
+ encoded_dict[encoded_key] = encoded_value
+ return encoded_dict
+ if isinstance(obj, (list, set, frozenset, GeneratorType, tuple)):
+ encoded_list = []
+ for item in obj:
+ encoded_list.append(jsonable_encoder(item, custom_encoder=custom_encoder))
+ return encoded_list
+
+ def fallback_serializer(o: Any) -> Any:
+ attempt_encode = encode_by_type(o)
+ if attempt_encode is not None:
+ return attempt_encode
+
+ try:
+ data = dict(o)
+ except Exception as e:
+ errors: List[Exception] = []
+ errors.append(e)
+ try:
+ data = vars(o)
+ except Exception as e:
+ errors.append(e)
+ raise ValueError(errors) from e
+ return jsonable_encoder(data, custom_encoder=custom_encoder)
+
+ return to_jsonable_with_fallback(obj, fallback_serializer)
diff --git a/v2/skyflow/generated/rest/core/pydantic_utilities.py b/v2/skyflow/generated/rest/core/pydantic_utilities.py
new file mode 100644
index 00000000..7db29500
--- /dev/null
+++ b/v2/skyflow/generated/rest/core/pydantic_utilities.py
@@ -0,0 +1,255 @@
+# This file was auto-generated by Fern from our API Definition.
+
+# nopycln: file
+import datetime as dt
+from collections import defaultdict
+from typing import Any, Callable, ClassVar, Dict, List, Mapping, Optional, Set, Tuple, Type, TypeVar, Union, cast
+
+import pydantic
+
+IS_PYDANTIC_V2 = pydantic.VERSION.startswith("2.")
+
+if IS_PYDANTIC_V2:
+ from pydantic.v1.datetime_parse import parse_date as parse_date
+ from pydantic.v1.datetime_parse import parse_datetime as parse_datetime
+ from pydantic.v1.fields import ModelField as ModelField
+ from pydantic.v1.json import ENCODERS_BY_TYPE as encoders_by_type # type: ignore[attr-defined]
+ from pydantic.v1.typing import get_args as get_args
+ from pydantic.v1.typing import get_origin as get_origin
+ from pydantic.v1.typing import is_literal_type as is_literal_type
+ from pydantic.v1.typing import is_union as is_union
+else:
+ from pydantic.datetime_parse import parse_date as parse_date # type: ignore[no-redef]
+ from pydantic.datetime_parse import parse_datetime as parse_datetime # type: ignore[no-redef]
+ from pydantic.fields import ModelField as ModelField # type: ignore[attr-defined, no-redef]
+ from pydantic.json import ENCODERS_BY_TYPE as encoders_by_type # type: ignore[no-redef]
+ from pydantic.typing import get_args as get_args # type: ignore[no-redef]
+ from pydantic.typing import get_origin as get_origin # type: ignore[no-redef]
+ from pydantic.typing import is_literal_type as is_literal_type # type: ignore[no-redef]
+ from pydantic.typing import is_union as is_union # type: ignore[no-redef]
+
+from .datetime_utils import serialize_datetime
+from .serialization import convert_and_respect_annotation_metadata
+from typing_extensions import TypeAlias
+
+T = TypeVar("T")
+Model = TypeVar("Model", bound=pydantic.BaseModel)
+
+
+def parse_obj_as(type_: Type[T], object_: Any) -> T:
+ dealiased_object = convert_and_respect_annotation_metadata(object_=object_, annotation=type_, direction="read")
+ if IS_PYDANTIC_V2:
+ adapter = pydantic.TypeAdapter(type_) # type: ignore[attr-defined]
+ return adapter.validate_python(dealiased_object)
+ return pydantic.parse_obj_as(type_, dealiased_object)
+
+
+def to_jsonable_with_fallback(obj: Any, fallback_serializer: Callable[[Any], Any]) -> Any:
+ if IS_PYDANTIC_V2:
+ from pydantic_core import to_jsonable_python
+
+ return to_jsonable_python(obj, fallback=fallback_serializer)
+ return fallback_serializer(obj)
+
+
+class UniversalBaseModel(pydantic.BaseModel):
+ if IS_PYDANTIC_V2:
+ model_config: ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict( # type: ignore[typeddict-unknown-key]
+ # Allow fields beginning with `model_` to be used in the model
+ protected_namespaces=(),
+ )
+
+ @pydantic.model_serializer(mode="plain", when_used="json") # type: ignore[attr-defined]
+ def serialize_model(self) -> Any: # type: ignore[name-defined]
+ serialized = self.model_dump()
+ data = {k: serialize_datetime(v) if isinstance(v, dt.datetime) else v for k, v in serialized.items()}
+ return data
+
+ else:
+
+ class Config:
+ smart_union = True
+ json_encoders = {dt.datetime: serialize_datetime}
+
+ @classmethod
+ def model_construct(cls: Type["Model"], _fields_set: Optional[Set[str]] = None, **values: Any) -> "Model":
+ dealiased_object = convert_and_respect_annotation_metadata(object_=values, annotation=cls, direction="read")
+ return cls.construct(_fields_set, **dealiased_object)
+
+ @classmethod
+ def construct(cls: Type["Model"], _fields_set: Optional[Set[str]] = None, **values: Any) -> "Model":
+ dealiased_object = convert_and_respect_annotation_metadata(object_=values, annotation=cls, direction="read")
+ if IS_PYDANTIC_V2:
+ return super().model_construct(_fields_set, **dealiased_object) # type: ignore[misc]
+ return super().construct(_fields_set, **dealiased_object)
+
+ def json(self, **kwargs: Any) -> str:
+ kwargs_with_defaults = {
+ "by_alias": True,
+ "exclude_unset": True,
+ **kwargs,
+ }
+ if IS_PYDANTIC_V2:
+ return super().model_dump_json(**kwargs_with_defaults) # type: ignore[misc]
+ return super().json(**kwargs_with_defaults)
+
+ def dict(self, **kwargs: Any) -> Dict[str, Any]:
+ """
+ Override the default dict method to `exclude_unset` by default. This function patches
+ `exclude_unset` to work include fields within non-None default values.
+ """
+ # Note: the logic here is multiplexed given the levers exposed in Pydantic V1 vs V2
+ # Pydantic V1's .dict can be extremely slow, so we do not want to call it twice.
+ #
+ # We'd ideally do the same for Pydantic V2, but it shells out to a library to serialize models
+ # that we have less control over, and this is less intrusive than custom serializers for now.
+ if IS_PYDANTIC_V2:
+ kwargs_with_defaults_exclude_unset = {
+ **kwargs,
+ "by_alias": True,
+ "exclude_unset": True,
+ "exclude_none": False,
+ }
+ kwargs_with_defaults_exclude_none = {
+ **kwargs,
+ "by_alias": True,
+ "exclude_none": True,
+ "exclude_unset": False,
+ }
+ dict_dump = deep_union_pydantic_dicts(
+ super().model_dump(**kwargs_with_defaults_exclude_unset), # type: ignore[misc]
+ super().model_dump(**kwargs_with_defaults_exclude_none), # type: ignore[misc]
+ )
+
+ else:
+ _fields_set = self.__fields_set__.copy()
+
+ fields = _get_model_fields(self.__class__)
+ for name, field in fields.items():
+ if name not in _fields_set:
+ default = _get_field_default(field)
+
+ # If the default values are non-null act like they've been set
+ # This effectively allows exclude_unset to work like exclude_none where
+ # the latter passes through intentionally set none values.
+ if default is not None or ("exclude_unset" in kwargs and not kwargs["exclude_unset"]):
+ _fields_set.add(name)
+
+ if default is not None:
+ self.__fields_set__.add(name)
+
+ kwargs_with_defaults_exclude_unset_include_fields = {
+ "by_alias": True,
+ "exclude_unset": True,
+ "include": _fields_set,
+ **kwargs,
+ }
+
+ dict_dump = super().dict(**kwargs_with_defaults_exclude_unset_include_fields)
+
+ return convert_and_respect_annotation_metadata(object_=dict_dump, annotation=self.__class__, direction="write")
+
+
+def _union_list_of_pydantic_dicts(source: List[Any], destination: List[Any]) -> List[Any]:
+ converted_list: List[Any] = []
+ for i, item in enumerate(source):
+ destination_value = destination[i]
+ if isinstance(item, dict):
+ converted_list.append(deep_union_pydantic_dicts(item, destination_value))
+ elif isinstance(item, list):
+ converted_list.append(_union_list_of_pydantic_dicts(item, destination_value))
+ else:
+ converted_list.append(item)
+ return converted_list
+
+
+def deep_union_pydantic_dicts(source: Dict[str, Any], destination: Dict[str, Any]) -> Dict[str, Any]:
+ for key, value in source.items():
+ node = destination.setdefault(key, {})
+ if isinstance(value, dict):
+ deep_union_pydantic_dicts(value, node)
+ # Note: we do not do this same processing for sets given we do not have sets of models
+ # and given the sets are unordered, the processing of the set and matching objects would
+ # be non-trivial.
+ elif isinstance(value, list):
+ destination[key] = _union_list_of_pydantic_dicts(value, node)
+ else:
+ destination[key] = value
+
+ return destination
+
+
+if IS_PYDANTIC_V2:
+
+ class V2RootModel(UniversalBaseModel, pydantic.RootModel): # type: ignore[misc, name-defined, type-arg]
+ pass
+
+ UniversalRootModel: TypeAlias = V2RootModel # type: ignore[misc]
+else:
+ UniversalRootModel: TypeAlias = UniversalBaseModel # type: ignore[misc, no-redef]
+
+
+def encode_by_type(o: Any) -> Any:
+ encoders_by_class_tuples: Dict[Callable[[Any], Any], Tuple[Any, ...]] = defaultdict(tuple)
+ for type_, encoder in encoders_by_type.items():
+ encoders_by_class_tuples[encoder] += (type_,)
+
+ if type(o) in encoders_by_type:
+ return encoders_by_type[type(o)](o)
+ for encoder, classes_tuple in encoders_by_class_tuples.items():
+ if isinstance(o, classes_tuple):
+ return encoder(o)
+
+
+def update_forward_refs(model: Type["Model"], **localns: Any) -> None:
+ if IS_PYDANTIC_V2:
+ model.model_rebuild(raise_errors=False) # type: ignore[attr-defined]
+ else:
+ model.update_forward_refs(**localns)
+
+
+# Mirrors Pydantic's internal typing
+AnyCallable = Callable[..., Any]
+
+
+def universal_root_validator(
+ pre: bool = False,
+) -> Callable[[AnyCallable], AnyCallable]:
+ def decorator(func: AnyCallable) -> AnyCallable:
+ if IS_PYDANTIC_V2:
+ return cast(AnyCallable, pydantic.model_validator(mode="before" if pre else "after")(func)) # type: ignore[attr-defined]
+ return cast(AnyCallable, pydantic.root_validator(pre=pre)(func)) # type: ignore[call-overload]
+
+ return decorator
+
+
+def universal_field_validator(field_name: str, pre: bool = False) -> Callable[[AnyCallable], AnyCallable]:
+ def decorator(func: AnyCallable) -> AnyCallable:
+ if IS_PYDANTIC_V2:
+ return cast(AnyCallable, pydantic.field_validator(field_name, mode="before" if pre else "after")(func)) # type: ignore[attr-defined]
+ return cast(AnyCallable, pydantic.validator(field_name, pre=pre)(func))
+
+ return decorator
+
+
+PydanticField = Union[ModelField, pydantic.fields.FieldInfo]
+
+
+def _get_model_fields(model: Type["Model"]) -> Mapping[str, PydanticField]:
+ if IS_PYDANTIC_V2:
+ return cast(Mapping[str, PydanticField], model.model_fields) # type: ignore[attr-defined]
+ return cast(Mapping[str, PydanticField], model.__fields__)
+
+
+def _get_field_default(field: PydanticField) -> Any:
+ try:
+ value = field.get_default() # type: ignore[union-attr]
+ except:
+ value = field.default
+ if IS_PYDANTIC_V2:
+ from pydantic_core import PydanticUndefined
+
+ if value == PydanticUndefined:
+ return None
+ return value
+ return value
diff --git a/v2/skyflow/generated/rest/core/query_encoder.py b/v2/skyflow/generated/rest/core/query_encoder.py
new file mode 100644
index 00000000..3183001d
--- /dev/null
+++ b/v2/skyflow/generated/rest/core/query_encoder.py
@@ -0,0 +1,58 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import pydantic
+
+
+# Flattens dicts to be of the form {"key[subkey][subkey2]": value} where value is not a dict
+def traverse_query_dict(dict_flat: Dict[str, Any], key_prefix: Optional[str] = None) -> List[Tuple[str, Any]]:
+ result = []
+ for k, v in dict_flat.items():
+ key = f"{key_prefix}[{k}]" if key_prefix is not None else k
+ if isinstance(v, dict):
+ result.extend(traverse_query_dict(v, key))
+ elif isinstance(v, list):
+ for arr_v in v:
+ if isinstance(arr_v, dict):
+ result.extend(traverse_query_dict(arr_v, key))
+ else:
+ result.append((key, arr_v))
+ else:
+ result.append((key, v))
+ return result
+
+
+def single_query_encoder(query_key: str, query_value: Any) -> List[Tuple[str, Any]]:
+ if isinstance(query_value, pydantic.BaseModel) or isinstance(query_value, dict):
+ if isinstance(query_value, pydantic.BaseModel):
+ obj_dict = query_value.dict(by_alias=True)
+ else:
+ obj_dict = query_value
+ return traverse_query_dict(obj_dict, query_key)
+ elif isinstance(query_value, list):
+ encoded_values: List[Tuple[str, Any]] = []
+ for value in query_value:
+ if isinstance(value, pydantic.BaseModel) or isinstance(value, dict):
+ if isinstance(value, pydantic.BaseModel):
+ obj_dict = value.dict(by_alias=True)
+ elif isinstance(value, dict):
+ obj_dict = value
+
+ encoded_values.extend(single_query_encoder(query_key, obj_dict))
+ else:
+ encoded_values.append((query_key, value))
+
+ return encoded_values
+
+ return [(query_key, query_value)]
+
+
+def encode_query(query: Optional[Dict[str, Any]]) -> Optional[List[Tuple[str, Any]]]:
+ if query is None:
+ return None
+
+ encoded_query = []
+ for k, v in query.items():
+ encoded_query.extend(single_query_encoder(k, v))
+ return encoded_query
diff --git a/v2/skyflow/generated/rest/core/remove_none_from_dict.py b/v2/skyflow/generated/rest/core/remove_none_from_dict.py
new file mode 100644
index 00000000..c2298143
--- /dev/null
+++ b/v2/skyflow/generated/rest/core/remove_none_from_dict.py
@@ -0,0 +1,11 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from typing import Any, Dict, Mapping, Optional
+
+
+def remove_none_from_dict(original: Mapping[str, Optional[Any]]) -> Dict[str, Any]:
+ new: Dict[str, Any] = {}
+ for key, value in original.items():
+ if value is not None:
+ new[key] = value
+ return new
diff --git a/v2/skyflow/generated/rest/core/request_options.py b/v2/skyflow/generated/rest/core/request_options.py
new file mode 100644
index 00000000..1b388044
--- /dev/null
+++ b/v2/skyflow/generated/rest/core/request_options.py
@@ -0,0 +1,35 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+try:
+ from typing import NotRequired # type: ignore
+except ImportError:
+ from typing_extensions import NotRequired
+
+
+class RequestOptions(typing.TypedDict, total=False):
+ """
+ Additional options for request-specific configuration when calling APIs via the SDK.
+ This is used primarily as an optional final parameter for service functions.
+
+ Attributes:
+ - timeout_in_seconds: int. The number of seconds to await an API call before timing out.
+
+ - max_retries: int. The max number of retries to attempt if the API call fails.
+
+ - additional_headers: typing.Dict[str, typing.Any]. A dictionary containing additional parameters to spread into the request's header dict
+
+ - additional_query_parameters: typing.Dict[str, typing.Any]. A dictionary containing additional parameters to spread into the request's query parameters dict
+
+ - additional_body_parameters: typing.Dict[str, typing.Any]. A dictionary containing additional parameters to spread into the request's body parameters dict
+
+ - chunk_size: int. The size, in bytes, to process each chunk of data being streamed back within the response. This equates to leveraging `chunk_size` within `requests` or `httpx`, and is only leveraged for file downloads.
+ """
+
+ timeout_in_seconds: NotRequired[int]
+ max_retries: NotRequired[int]
+ additional_headers: NotRequired[typing.Dict[str, typing.Any]]
+ additional_query_parameters: NotRequired[typing.Dict[str, typing.Any]]
+ additional_body_parameters: NotRequired[typing.Dict[str, typing.Any]]
+ chunk_size: NotRequired[int]
diff --git a/v2/skyflow/generated/rest/core/serialization.py b/v2/skyflow/generated/rest/core/serialization.py
new file mode 100644
index 00000000..c36e865c
--- /dev/null
+++ b/v2/skyflow/generated/rest/core/serialization.py
@@ -0,0 +1,276 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import collections
+import inspect
+import typing
+
+import pydantic
+import typing_extensions
+
+
+class FieldMetadata:
+ """
+ Metadata class used to annotate fields to provide additional information.
+
+ Example:
+ class MyDict(TypedDict):
+ field: typing.Annotated[str, FieldMetadata(alias="field_name")]
+
+ Will serialize: `{"field": "value"}`
+ To: `{"field_name": "value"}`
+ """
+
+ alias: str
+
+ def __init__(self, *, alias: str) -> None:
+ self.alias = alias
+
+
+def convert_and_respect_annotation_metadata(
+ *,
+ object_: typing.Any,
+ annotation: typing.Any,
+ inner_type: typing.Optional[typing.Any] = None,
+ direction: typing.Literal["read", "write"],
+) -> typing.Any:
+ """
+ Respect the metadata annotations on a field, such as aliasing. This function effectively
+ manipulates the dict-form of an object to respect the metadata annotations. This is primarily used for
+ TypedDicts, which cannot support aliasing out of the box, and can be extended for additional
+ utilities, such as defaults.
+
+ Parameters
+ ----------
+ object_ : typing.Any
+
+ annotation : type
+ The type we're looking to apply typing annotations from
+
+ inner_type : typing.Optional[type]
+
+ Returns
+ -------
+ typing.Any
+ """
+
+ if object_ is None:
+ return None
+ if inner_type is None:
+ inner_type = annotation
+
+ clean_type = _remove_annotations(inner_type)
+ # Pydantic models
+ if (
+ inspect.isclass(clean_type)
+ and issubclass(clean_type, pydantic.BaseModel)
+ and isinstance(object_, typing.Mapping)
+ ):
+ return _convert_mapping(object_, clean_type, direction)
+ # TypedDicts
+ if typing_extensions.is_typeddict(clean_type) and isinstance(object_, typing.Mapping):
+ return _convert_mapping(object_, clean_type, direction)
+
+ if (
+ typing_extensions.get_origin(clean_type) == typing.Dict
+ or typing_extensions.get_origin(clean_type) == dict
+ or clean_type == typing.Dict
+ ) and isinstance(object_, typing.Dict):
+ key_type = typing_extensions.get_args(clean_type)[0]
+ value_type = typing_extensions.get_args(clean_type)[1]
+
+ return {
+ key: convert_and_respect_annotation_metadata(
+ object_=value,
+ annotation=annotation,
+ inner_type=value_type,
+ direction=direction,
+ )
+ for key, value in object_.items()
+ }
+
+ # If you're iterating on a string, do not bother to coerce it to a sequence.
+ if not isinstance(object_, str):
+ if (
+ typing_extensions.get_origin(clean_type) == typing.Set
+ or typing_extensions.get_origin(clean_type) == set
+ or clean_type == typing.Set
+ ) and isinstance(object_, typing.Set):
+ inner_type = typing_extensions.get_args(clean_type)[0]
+ return {
+ convert_and_respect_annotation_metadata(
+ object_=item,
+ annotation=annotation,
+ inner_type=inner_type,
+ direction=direction,
+ )
+ for item in object_
+ }
+ elif (
+ (
+ typing_extensions.get_origin(clean_type) == typing.List
+ or typing_extensions.get_origin(clean_type) == list
+ or clean_type == typing.List
+ )
+ and isinstance(object_, typing.List)
+ ) or (
+ (
+ typing_extensions.get_origin(clean_type) == typing.Sequence
+ or typing_extensions.get_origin(clean_type) == collections.abc.Sequence
+ or clean_type == typing.Sequence
+ )
+ and isinstance(object_, typing.Sequence)
+ ):
+ inner_type = typing_extensions.get_args(clean_type)[0]
+ return [
+ convert_and_respect_annotation_metadata(
+ object_=item,
+ annotation=annotation,
+ inner_type=inner_type,
+ direction=direction,
+ )
+ for item in object_
+ ]
+
+ if typing_extensions.get_origin(clean_type) == typing.Union:
+ # We should be able to ~relatively~ safely try to convert keys against all
+ # member types in the union, the edge case here is if one member aliases a field
+ # of the same name to a different name from another member
+ # Or if another member aliases a field of the same name that another member does not.
+ for member in typing_extensions.get_args(clean_type):
+ object_ = convert_and_respect_annotation_metadata(
+ object_=object_,
+ annotation=annotation,
+ inner_type=member,
+ direction=direction,
+ )
+ return object_
+
+ annotated_type = _get_annotation(annotation)
+ if annotated_type is None:
+ return object_
+
+ # If the object is not a TypedDict, a Union, or other container (list, set, sequence, etc.)
+ # Then we can safely call it on the recursive conversion.
+ return object_
+
+
+def _convert_mapping(
+ object_: typing.Mapping[str, object],
+ expected_type: typing.Any,
+ direction: typing.Literal["read", "write"],
+) -> typing.Mapping[str, object]:
+ converted_object: typing.Dict[str, object] = {}
+ try:
+ annotations = typing_extensions.get_type_hints(expected_type, include_extras=True)
+ except NameError:
+ # The TypedDict contains a circular reference, so
+ # we use the __annotations__ attribute directly.
+ annotations = getattr(expected_type, "__annotations__", {})
+ aliases_to_field_names = _get_alias_to_field_name(annotations)
+ for key, value in object_.items():
+ if direction == "read" and key in aliases_to_field_names:
+ dealiased_key = aliases_to_field_names.get(key)
+ if dealiased_key is not None:
+ type_ = annotations.get(dealiased_key)
+ else:
+ type_ = annotations.get(key)
+ # Note you can't get the annotation by the field name if you're in read mode, so you must check the aliases map
+ #
+ # So this is effectively saying if we're in write mode, and we don't have a type, or if we're in read mode and we don't have an alias
+ # then we can just pass the value through as is
+ if type_ is None:
+ converted_object[key] = value
+ elif direction == "read" and key not in aliases_to_field_names:
+ converted_object[key] = convert_and_respect_annotation_metadata(
+ object_=value, annotation=type_, direction=direction
+ )
+ else:
+ converted_object[_alias_key(key, type_, direction, aliases_to_field_names)] = (
+ convert_and_respect_annotation_metadata(object_=value, annotation=type_, direction=direction)
+ )
+ return converted_object
+
+
+def _get_annotation(type_: typing.Any) -> typing.Optional[typing.Any]:
+ maybe_annotated_type = typing_extensions.get_origin(type_)
+ if maybe_annotated_type is None:
+ return None
+
+ if maybe_annotated_type == typing_extensions.NotRequired:
+ type_ = typing_extensions.get_args(type_)[0]
+ maybe_annotated_type = typing_extensions.get_origin(type_)
+
+ if maybe_annotated_type == typing_extensions.Annotated:
+ return type_
+
+ return None
+
+
+def _remove_annotations(type_: typing.Any) -> typing.Any:
+ maybe_annotated_type = typing_extensions.get_origin(type_)
+ if maybe_annotated_type is None:
+ return type_
+
+ if maybe_annotated_type == typing_extensions.NotRequired:
+ return _remove_annotations(typing_extensions.get_args(type_)[0])
+
+ if maybe_annotated_type == typing_extensions.Annotated:
+ return _remove_annotations(typing_extensions.get_args(type_)[0])
+
+ return type_
+
+
+def get_alias_to_field_mapping(type_: typing.Any) -> typing.Dict[str, str]:
+ annotations = typing_extensions.get_type_hints(type_, include_extras=True)
+ return _get_alias_to_field_name(annotations)
+
+
+def get_field_to_alias_mapping(type_: typing.Any) -> typing.Dict[str, str]:
+ annotations = typing_extensions.get_type_hints(type_, include_extras=True)
+ return _get_field_to_alias_name(annotations)
+
+
+def _get_alias_to_field_name(
+ field_to_hint: typing.Dict[str, typing.Any],
+) -> typing.Dict[str, str]:
+ aliases = {}
+ for field, hint in field_to_hint.items():
+ maybe_alias = _get_alias_from_type(hint)
+ if maybe_alias is not None:
+ aliases[maybe_alias] = field
+ return aliases
+
+
+def _get_field_to_alias_name(
+ field_to_hint: typing.Dict[str, typing.Any],
+) -> typing.Dict[str, str]:
+ aliases = {}
+ for field, hint in field_to_hint.items():
+ maybe_alias = _get_alias_from_type(hint)
+ if maybe_alias is not None:
+ aliases[field] = maybe_alias
+ return aliases
+
+
+def _get_alias_from_type(type_: typing.Any) -> typing.Optional[str]:
+ maybe_annotated_type = _get_annotation(type_)
+
+ if maybe_annotated_type is not None:
+ # The actual annotations are 1 onward, the first is the annotated type
+ annotations = typing_extensions.get_args(maybe_annotated_type)[1:]
+
+ for annotation in annotations:
+ if isinstance(annotation, FieldMetadata) and annotation.alias is not None:
+ return annotation.alias
+ return None
+
+
+def _alias_key(
+ key: str,
+ type_: typing.Any,
+ direction: typing.Literal["read", "write"],
+ aliases_to_field_names: typing.Dict[str, str],
+) -> str:
+ if direction == "read":
+ return aliases_to_field_names.get(key, key)
+ return _get_alias_from_type(type_=type_) or key
diff --git a/skyflow/generated/rest/environment.py b/v2/skyflow/generated/rest/environment.py
similarity index 100%
rename from skyflow/generated/rest/environment.py
rename to v2/skyflow/generated/rest/environment.py
diff --git a/skyflow/generated/rest/errors/__init__.py b/v2/skyflow/generated/rest/errors/__init__.py
similarity index 100%
rename from skyflow/generated/rest/errors/__init__.py
rename to v2/skyflow/generated/rest/errors/__init__.py
diff --git a/skyflow/generated/rest/errors/bad_request_error.py b/v2/skyflow/generated/rest/errors/bad_request_error.py
similarity index 100%
rename from skyflow/generated/rest/errors/bad_request_error.py
rename to v2/skyflow/generated/rest/errors/bad_request_error.py
diff --git a/skyflow/generated/rest/errors/internal_server_error.py b/v2/skyflow/generated/rest/errors/internal_server_error.py
similarity index 100%
rename from skyflow/generated/rest/errors/internal_server_error.py
rename to v2/skyflow/generated/rest/errors/internal_server_error.py
diff --git a/skyflow/generated/rest/errors/not_found_error.py b/v2/skyflow/generated/rest/errors/not_found_error.py
similarity index 100%
rename from skyflow/generated/rest/errors/not_found_error.py
rename to v2/skyflow/generated/rest/errors/not_found_error.py
diff --git a/skyflow/generated/rest/errors/unauthorized_error.py b/v2/skyflow/generated/rest/errors/unauthorized_error.py
similarity index 100%
rename from skyflow/generated/rest/errors/unauthorized_error.py
rename to v2/skyflow/generated/rest/errors/unauthorized_error.py
diff --git a/skyflow/generated/rest/files/__init__.py b/v2/skyflow/generated/rest/files/__init__.py
similarity index 100%
rename from skyflow/generated/rest/files/__init__.py
rename to v2/skyflow/generated/rest/files/__init__.py
diff --git a/skyflow/generated/rest/files/client.py b/v2/skyflow/generated/rest/files/client.py
similarity index 100%
rename from skyflow/generated/rest/files/client.py
rename to v2/skyflow/generated/rest/files/client.py
diff --git a/skyflow/generated/rest/files/raw_client.py b/v2/skyflow/generated/rest/files/raw_client.py
similarity index 100%
rename from skyflow/generated/rest/files/raw_client.py
rename to v2/skyflow/generated/rest/files/raw_client.py
diff --git a/skyflow/generated/rest/files/types/__init__.py b/v2/skyflow/generated/rest/files/types/__init__.py
similarity index 100%
rename from skyflow/generated/rest/files/types/__init__.py
rename to v2/skyflow/generated/rest/files/types/__init__.py
diff --git a/skyflow/generated/rest/files/types/deidentify_file_audio_request_deidentify_audio_entity_types_item.py b/v2/skyflow/generated/rest/files/types/deidentify_file_audio_request_deidentify_audio_entity_types_item.py
similarity index 100%
rename from skyflow/generated/rest/files/types/deidentify_file_audio_request_deidentify_audio_entity_types_item.py
rename to v2/skyflow/generated/rest/files/types/deidentify_file_audio_request_deidentify_audio_entity_types_item.py
diff --git a/skyflow/generated/rest/files/types/deidentify_file_audio_request_deidentify_audio_output_transcription.py b/v2/skyflow/generated/rest/files/types/deidentify_file_audio_request_deidentify_audio_output_transcription.py
similarity index 100%
rename from skyflow/generated/rest/files/types/deidentify_file_audio_request_deidentify_audio_output_transcription.py
rename to v2/skyflow/generated/rest/files/types/deidentify_file_audio_request_deidentify_audio_output_transcription.py
diff --git a/skyflow/generated/rest/files/types/deidentify_file_document_pdf_request_deidentify_pdf_entity_types_item.py b/v2/skyflow/generated/rest/files/types/deidentify_file_document_pdf_request_deidentify_pdf_entity_types_item.py
similarity index 100%
rename from skyflow/generated/rest/files/types/deidentify_file_document_pdf_request_deidentify_pdf_entity_types_item.py
rename to v2/skyflow/generated/rest/files/types/deidentify_file_document_pdf_request_deidentify_pdf_entity_types_item.py
diff --git a/skyflow/generated/rest/files/types/deidentify_file_image_request_deidentify_image_entity_types_item.py b/v2/skyflow/generated/rest/files/types/deidentify_file_image_request_deidentify_image_entity_types_item.py
similarity index 100%
rename from skyflow/generated/rest/files/types/deidentify_file_image_request_deidentify_image_entity_types_item.py
rename to v2/skyflow/generated/rest/files/types/deidentify_file_image_request_deidentify_image_entity_types_item.py
diff --git a/skyflow/generated/rest/files/types/deidentify_file_image_request_deidentify_image_masking_method.py b/v2/skyflow/generated/rest/files/types/deidentify_file_image_request_deidentify_image_masking_method.py
similarity index 100%
rename from skyflow/generated/rest/files/types/deidentify_file_image_request_deidentify_image_masking_method.py
rename to v2/skyflow/generated/rest/files/types/deidentify_file_image_request_deidentify_image_masking_method.py
diff --git a/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_document_entity_types_item.py b/v2/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_document_entity_types_item.py
similarity index 100%
rename from skyflow/generated/rest/files/types/deidentify_file_request_deidentify_document_entity_types_item.py
rename to v2/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_document_entity_types_item.py
diff --git a/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_presentation_entity_types_item.py b/v2/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_presentation_entity_types_item.py
similarity index 100%
rename from skyflow/generated/rest/files/types/deidentify_file_request_deidentify_presentation_entity_types_item.py
rename to v2/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_presentation_entity_types_item.py
diff --git a/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_spreadsheet_entity_types_item.py b/v2/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_spreadsheet_entity_types_item.py
similarity index 100%
rename from skyflow/generated/rest/files/types/deidentify_file_request_deidentify_spreadsheet_entity_types_item.py
rename to v2/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_spreadsheet_entity_types_item.py
diff --git a/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_structured_text_entity_types_item.py b/v2/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_structured_text_entity_types_item.py
similarity index 100%
rename from skyflow/generated/rest/files/types/deidentify_file_request_deidentify_structured_text_entity_types_item.py
rename to v2/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_structured_text_entity_types_item.py
diff --git a/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_text_entity_types_item.py b/v2/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_text_entity_types_item.py
similarity index 100%
rename from skyflow/generated/rest/files/types/deidentify_file_request_deidentify_text_entity_types_item.py
rename to v2/skyflow/generated/rest/files/types/deidentify_file_request_deidentify_text_entity_types_item.py
diff --git a/skyflow/generated/rest/files/types/deidentify_file_request_entity_types_item.py b/v2/skyflow/generated/rest/files/types/deidentify_file_request_entity_types_item.py
similarity index 100%
rename from skyflow/generated/rest/files/types/deidentify_file_request_entity_types_item.py
rename to v2/skyflow/generated/rest/files/types/deidentify_file_request_entity_types_item.py
diff --git a/v2/skyflow/generated/rest/guardrails/__init__.py b/v2/skyflow/generated/rest/guardrails/__init__.py
new file mode 100644
index 00000000..5cde0202
--- /dev/null
+++ b/v2/skyflow/generated/rest/guardrails/__init__.py
@@ -0,0 +1,4 @@
+# This file was auto-generated by Fern from our API Definition.
+
+# isort: skip_file
+
diff --git a/skyflow/generated/rest/guardrails/client.py b/v2/skyflow/generated/rest/guardrails/client.py
similarity index 100%
rename from skyflow/generated/rest/guardrails/client.py
rename to v2/skyflow/generated/rest/guardrails/client.py
diff --git a/skyflow/generated/rest/guardrails/raw_client.py b/v2/skyflow/generated/rest/guardrails/raw_client.py
similarity index 100%
rename from skyflow/generated/rest/guardrails/raw_client.py
rename to v2/skyflow/generated/rest/guardrails/raw_client.py
diff --git a/v2/skyflow/generated/rest/py.typed b/v2/skyflow/generated/rest/py.typed
new file mode 100644
index 00000000..e69de29b
diff --git a/v2/skyflow/generated/rest/query/__init__.py b/v2/skyflow/generated/rest/query/__init__.py
new file mode 100644
index 00000000..5cde0202
--- /dev/null
+++ b/v2/skyflow/generated/rest/query/__init__.py
@@ -0,0 +1,4 @@
+# This file was auto-generated by Fern from our API Definition.
+
+# isort: skip_file
+
diff --git a/skyflow/generated/rest/query/client.py b/v2/skyflow/generated/rest/query/client.py
similarity index 100%
rename from skyflow/generated/rest/query/client.py
rename to v2/skyflow/generated/rest/query/client.py
diff --git a/skyflow/generated/rest/query/raw_client.py b/v2/skyflow/generated/rest/query/raw_client.py
similarity index 100%
rename from skyflow/generated/rest/query/raw_client.py
rename to v2/skyflow/generated/rest/query/raw_client.py
diff --git a/skyflow/generated/rest/records/__init__.py b/v2/skyflow/generated/rest/records/__init__.py
similarity index 100%
rename from skyflow/generated/rest/records/__init__.py
rename to v2/skyflow/generated/rest/records/__init__.py
diff --git a/skyflow/generated/rest/records/client.py b/v2/skyflow/generated/rest/records/client.py
similarity index 100%
rename from skyflow/generated/rest/records/client.py
rename to v2/skyflow/generated/rest/records/client.py
diff --git a/skyflow/generated/rest/records/raw_client.py b/v2/skyflow/generated/rest/records/raw_client.py
similarity index 100%
rename from skyflow/generated/rest/records/raw_client.py
rename to v2/skyflow/generated/rest/records/raw_client.py
diff --git a/skyflow/generated/rest/records/types/__init__.py b/v2/skyflow/generated/rest/records/types/__init__.py
similarity index 100%
rename from skyflow/generated/rest/records/types/__init__.py
rename to v2/skyflow/generated/rest/records/types/__init__.py
diff --git a/skyflow/generated/rest/records/types/record_service_bulk_get_record_request_order_by.py b/v2/skyflow/generated/rest/records/types/record_service_bulk_get_record_request_order_by.py
similarity index 100%
rename from skyflow/generated/rest/records/types/record_service_bulk_get_record_request_order_by.py
rename to v2/skyflow/generated/rest/records/types/record_service_bulk_get_record_request_order_by.py
diff --git a/skyflow/generated/rest/records/types/record_service_bulk_get_record_request_redaction.py b/v2/skyflow/generated/rest/records/types/record_service_bulk_get_record_request_redaction.py
similarity index 100%
rename from skyflow/generated/rest/records/types/record_service_bulk_get_record_request_redaction.py
rename to v2/skyflow/generated/rest/records/types/record_service_bulk_get_record_request_redaction.py
diff --git a/skyflow/generated/rest/records/types/record_service_get_record_request_redaction.py b/v2/skyflow/generated/rest/records/types/record_service_get_record_request_redaction.py
similarity index 100%
rename from skyflow/generated/rest/records/types/record_service_get_record_request_redaction.py
rename to v2/skyflow/generated/rest/records/types/record_service_get_record_request_redaction.py
diff --git a/skyflow/generated/rest/strings/__init__.py b/v2/skyflow/generated/rest/strings/__init__.py
similarity index 100%
rename from skyflow/generated/rest/strings/__init__.py
rename to v2/skyflow/generated/rest/strings/__init__.py
diff --git a/skyflow/generated/rest/strings/client.py b/v2/skyflow/generated/rest/strings/client.py
similarity index 100%
rename from skyflow/generated/rest/strings/client.py
rename to v2/skyflow/generated/rest/strings/client.py
diff --git a/skyflow/generated/rest/strings/raw_client.py b/v2/skyflow/generated/rest/strings/raw_client.py
similarity index 100%
rename from skyflow/generated/rest/strings/raw_client.py
rename to v2/skyflow/generated/rest/strings/raw_client.py
diff --git a/skyflow/generated/rest/strings/types/__init__.py b/v2/skyflow/generated/rest/strings/types/__init__.py
similarity index 100%
rename from skyflow/generated/rest/strings/types/__init__.py
rename to v2/skyflow/generated/rest/strings/types/__init__.py
diff --git a/skyflow/generated/rest/strings/types/deidentify_string_request_entity_types_item.py b/v2/skyflow/generated/rest/strings/types/deidentify_string_request_entity_types_item.py
similarity index 100%
rename from skyflow/generated/rest/strings/types/deidentify_string_request_entity_types_item.py
rename to v2/skyflow/generated/rest/strings/types/deidentify_string_request_entity_types_item.py
diff --git a/v2/skyflow/generated/rest/tokens/__init__.py b/v2/skyflow/generated/rest/tokens/__init__.py
new file mode 100644
index 00000000..5cde0202
--- /dev/null
+++ b/v2/skyflow/generated/rest/tokens/__init__.py
@@ -0,0 +1,4 @@
+# This file was auto-generated by Fern from our API Definition.
+
+# isort: skip_file
+
diff --git a/skyflow/generated/rest/tokens/client.py b/v2/skyflow/generated/rest/tokens/client.py
similarity index 100%
rename from skyflow/generated/rest/tokens/client.py
rename to v2/skyflow/generated/rest/tokens/client.py
diff --git a/skyflow/generated/rest/tokens/raw_client.py b/v2/skyflow/generated/rest/tokens/raw_client.py
similarity index 100%
rename from skyflow/generated/rest/tokens/raw_client.py
rename to v2/skyflow/generated/rest/tokens/raw_client.py
diff --git a/skyflow/generated/rest/types/__init__.py b/v2/skyflow/generated/rest/types/__init__.py
similarity index 100%
rename from skyflow/generated/rest/types/__init__.py
rename to v2/skyflow/generated/rest/types/__init__.py
diff --git a/skyflow/generated/rest/types/audit_event_audit_resource_type.py b/v2/skyflow/generated/rest/types/audit_event_audit_resource_type.py
similarity index 100%
rename from skyflow/generated/rest/types/audit_event_audit_resource_type.py
rename to v2/skyflow/generated/rest/types/audit_event_audit_resource_type.py
diff --git a/skyflow/generated/rest/types/audit_event_context.py b/v2/skyflow/generated/rest/types/audit_event_context.py
similarity index 100%
rename from skyflow/generated/rest/types/audit_event_context.py
rename to v2/skyflow/generated/rest/types/audit_event_context.py
diff --git a/skyflow/generated/rest/types/audit_event_data.py b/v2/skyflow/generated/rest/types/audit_event_data.py
similarity index 100%
rename from skyflow/generated/rest/types/audit_event_data.py
rename to v2/skyflow/generated/rest/types/audit_event_data.py
diff --git a/skyflow/generated/rest/types/audit_event_http_info.py b/v2/skyflow/generated/rest/types/audit_event_http_info.py
similarity index 100%
rename from skyflow/generated/rest/types/audit_event_http_info.py
rename to v2/skyflow/generated/rest/types/audit_event_http_info.py
diff --git a/skyflow/generated/rest/types/batch_record_method.py b/v2/skyflow/generated/rest/types/batch_record_method.py
similarity index 100%
rename from skyflow/generated/rest/types/batch_record_method.py
rename to v2/skyflow/generated/rest/types/batch_record_method.py
diff --git a/skyflow/generated/rest/types/context_access_type.py b/v2/skyflow/generated/rest/types/context_access_type.py
similarity index 100%
rename from skyflow/generated/rest/types/context_access_type.py
rename to v2/skyflow/generated/rest/types/context_access_type.py
diff --git a/skyflow/generated/rest/types/context_auth_mode.py b/v2/skyflow/generated/rest/types/context_auth_mode.py
similarity index 100%
rename from skyflow/generated/rest/types/context_auth_mode.py
rename to v2/skyflow/generated/rest/types/context_auth_mode.py
diff --git a/skyflow/generated/rest/types/deidentified_file_output.py b/v2/skyflow/generated/rest/types/deidentified_file_output.py
similarity index 100%
rename from skyflow/generated/rest/types/deidentified_file_output.py
rename to v2/skyflow/generated/rest/types/deidentified_file_output.py
diff --git a/skyflow/generated/rest/types/deidentified_file_output_processed_file_extension.py b/v2/skyflow/generated/rest/types/deidentified_file_output_processed_file_extension.py
similarity index 100%
rename from skyflow/generated/rest/types/deidentified_file_output_processed_file_extension.py
rename to v2/skyflow/generated/rest/types/deidentified_file_output_processed_file_extension.py
diff --git a/skyflow/generated/rest/types/deidentified_file_output_processed_file_type.py b/v2/skyflow/generated/rest/types/deidentified_file_output_processed_file_type.py
similarity index 100%
rename from skyflow/generated/rest/types/deidentified_file_output_processed_file_type.py
rename to v2/skyflow/generated/rest/types/deidentified_file_output_processed_file_type.py
diff --git a/skyflow/generated/rest/types/deidentify_file_response.py b/v2/skyflow/generated/rest/types/deidentify_file_response.py
similarity index 100%
rename from skyflow/generated/rest/types/deidentify_file_response.py
rename to v2/skyflow/generated/rest/types/deidentify_file_response.py
diff --git a/skyflow/generated/rest/types/deidentify_string_response.py b/v2/skyflow/generated/rest/types/deidentify_string_response.py
similarity index 100%
rename from skyflow/generated/rest/types/deidentify_string_response.py
rename to v2/skyflow/generated/rest/types/deidentify_string_response.py
diff --git a/skyflow/generated/rest/types/detect_guardrails_response.py b/v2/skyflow/generated/rest/types/detect_guardrails_response.py
similarity index 100%
rename from skyflow/generated/rest/types/detect_guardrails_response.py
rename to v2/skyflow/generated/rest/types/detect_guardrails_response.py
diff --git a/skyflow/generated/rest/types/detect_guardrails_response_validation.py b/v2/skyflow/generated/rest/types/detect_guardrails_response_validation.py
similarity index 100%
rename from skyflow/generated/rest/types/detect_guardrails_response_validation.py
rename to v2/skyflow/generated/rest/types/detect_guardrails_response_validation.py
diff --git a/skyflow/generated/rest/types/detect_runs_response.py b/v2/skyflow/generated/rest/types/detect_runs_response.py
similarity index 100%
rename from skyflow/generated/rest/types/detect_runs_response.py
rename to v2/skyflow/generated/rest/types/detect_runs_response.py
diff --git a/skyflow/generated/rest/types/detect_runs_response_output_type.py b/v2/skyflow/generated/rest/types/detect_runs_response_output_type.py
similarity index 100%
rename from skyflow/generated/rest/types/detect_runs_response_output_type.py
rename to v2/skyflow/generated/rest/types/detect_runs_response_output_type.py
diff --git a/skyflow/generated/rest/types/detect_runs_response_status.py b/v2/skyflow/generated/rest/types/detect_runs_response_status.py
similarity index 100%
rename from skyflow/generated/rest/types/detect_runs_response_status.py
rename to v2/skyflow/generated/rest/types/detect_runs_response_status.py
diff --git a/skyflow/generated/rest/types/detokenize_record_response_value_type.py b/v2/skyflow/generated/rest/types/detokenize_record_response_value_type.py
similarity index 100%
rename from skyflow/generated/rest/types/detokenize_record_response_value_type.py
rename to v2/skyflow/generated/rest/types/detokenize_record_response_value_type.py
diff --git a/skyflow/generated/rest/types/error_response.py b/v2/skyflow/generated/rest/types/error_response.py
similarity index 100%
rename from skyflow/generated/rest/types/error_response.py
rename to v2/skyflow/generated/rest/types/error_response.py
diff --git a/skyflow/generated/rest/types/error_response_error.py b/v2/skyflow/generated/rest/types/error_response_error.py
similarity index 100%
rename from skyflow/generated/rest/types/error_response_error.py
rename to v2/skyflow/generated/rest/types/error_response_error.py
diff --git a/skyflow/generated/rest/types/file_data.py b/v2/skyflow/generated/rest/types/file_data.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data.py
rename to v2/skyflow/generated/rest/types/file_data.py
diff --git a/skyflow/generated/rest/types/file_data_data_format.py b/v2/skyflow/generated/rest/types/file_data_data_format.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_data_format.py
rename to v2/skyflow/generated/rest/types/file_data_data_format.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_audio.py b/v2/skyflow/generated/rest/types/file_data_deidentify_audio.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_audio.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_audio.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_audio_data_format.py b/v2/skyflow/generated/rest/types/file_data_deidentify_audio_data_format.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_audio_data_format.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_audio_data_format.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_document.py b/v2/skyflow/generated/rest/types/file_data_deidentify_document.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_document.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_document.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_document_data_format.py b/v2/skyflow/generated/rest/types/file_data_deidentify_document_data_format.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_document_data_format.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_document_data_format.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_image.py b/v2/skyflow/generated/rest/types/file_data_deidentify_image.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_image.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_image.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_image_data_format.py b/v2/skyflow/generated/rest/types/file_data_deidentify_image_data_format.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_image_data_format.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_image_data_format.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_pdf.py b/v2/skyflow/generated/rest/types/file_data_deidentify_pdf.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_pdf.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_pdf.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_presentation.py b/v2/skyflow/generated/rest/types/file_data_deidentify_presentation.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_presentation.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_presentation.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_presentation_data_format.py b/v2/skyflow/generated/rest/types/file_data_deidentify_presentation_data_format.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_presentation_data_format.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_presentation_data_format.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_spreadsheet.py b/v2/skyflow/generated/rest/types/file_data_deidentify_spreadsheet.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_spreadsheet.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_spreadsheet.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_spreadsheet_data_format.py b/v2/skyflow/generated/rest/types/file_data_deidentify_spreadsheet_data_format.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_spreadsheet_data_format.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_spreadsheet_data_format.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_structured_text.py b/v2/skyflow/generated/rest/types/file_data_deidentify_structured_text.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_structured_text.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_structured_text.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_structured_text_data_format.py b/v2/skyflow/generated/rest/types/file_data_deidentify_structured_text_data_format.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_structured_text_data_format.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_structured_text_data_format.py
diff --git a/skyflow/generated/rest/types/file_data_deidentify_text.py b/v2/skyflow/generated/rest/types/file_data_deidentify_text.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_deidentify_text.py
rename to v2/skyflow/generated/rest/types/file_data_deidentify_text.py
diff --git a/skyflow/generated/rest/types/file_data_reidentify_file.py b/v2/skyflow/generated/rest/types/file_data_reidentify_file.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_reidentify_file.py
rename to v2/skyflow/generated/rest/types/file_data_reidentify_file.py
diff --git a/skyflow/generated/rest/types/file_data_reidentify_file_data_format.py b/v2/skyflow/generated/rest/types/file_data_reidentify_file_data_format.py
similarity index 100%
rename from skyflow/generated/rest/types/file_data_reidentify_file_data_format.py
rename to v2/skyflow/generated/rest/types/file_data_reidentify_file_data_format.py
diff --git a/skyflow/generated/rest/types/format.py b/v2/skyflow/generated/rest/types/format.py
similarity index 100%
rename from skyflow/generated/rest/types/format.py
rename to v2/skyflow/generated/rest/types/format.py
diff --git a/skyflow/generated/rest/types/format_masked_item.py b/v2/skyflow/generated/rest/types/format_masked_item.py
similarity index 100%
rename from skyflow/generated/rest/types/format_masked_item.py
rename to v2/skyflow/generated/rest/types/format_masked_item.py
diff --git a/skyflow/generated/rest/types/format_plaintext_item.py b/v2/skyflow/generated/rest/types/format_plaintext_item.py
similarity index 100%
rename from skyflow/generated/rest/types/format_plaintext_item.py
rename to v2/skyflow/generated/rest/types/format_plaintext_item.py
diff --git a/skyflow/generated/rest/types/format_redacted_item.py b/v2/skyflow/generated/rest/types/format_redacted_item.py
similarity index 100%
rename from skyflow/generated/rest/types/format_redacted_item.py
rename to v2/skyflow/generated/rest/types/format_redacted_item.py
diff --git a/v2/skyflow/generated/rest/types/googlerpc_status.py b/v2/skyflow/generated/rest/types/googlerpc_status.py
new file mode 100644
index 00000000..f0a885b4
--- /dev/null
+++ b/v2/skyflow/generated/rest/types/googlerpc_status.py
@@ -0,0 +1,22 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from .protobuf_any import ProtobufAny
+
+
+class GooglerpcStatus(UniversalBaseModel):
+ code: typing.Optional[int] = None
+ message: typing.Optional[str] = None
+ details: typing.Optional[typing.List[ProtobufAny]] = None
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/skyflow/generated/rest/types/http_code.py b/v2/skyflow/generated/rest/types/http_code.py
similarity index 100%
rename from skyflow/generated/rest/types/http_code.py
rename to v2/skyflow/generated/rest/types/http_code.py
diff --git a/skyflow/generated/rest/types/identify_response.py b/v2/skyflow/generated/rest/types/identify_response.py
similarity index 100%
rename from skyflow/generated/rest/types/identify_response.py
rename to v2/skyflow/generated/rest/types/identify_response.py
diff --git a/skyflow/generated/rest/types/locations.py b/v2/skyflow/generated/rest/types/locations.py
similarity index 100%
rename from skyflow/generated/rest/types/locations.py
rename to v2/skyflow/generated/rest/types/locations.py
diff --git a/v2/skyflow/generated/rest/types/protobuf_any.py b/v2/skyflow/generated/rest/types/protobuf_any.py
new file mode 100644
index 00000000..9062870c
--- /dev/null
+++ b/v2/skyflow/generated/rest/types/protobuf_any.py
@@ -0,0 +1,21 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class ProtobufAny(UniversalBaseModel):
+ type: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="@type")] = None
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/skyflow/generated/rest/types/redaction_enum_redaction.py b/v2/skyflow/generated/rest/types/redaction_enum_redaction.py
similarity index 100%
rename from skyflow/generated/rest/types/redaction_enum_redaction.py
rename to v2/skyflow/generated/rest/types/redaction_enum_redaction.py
diff --git a/skyflow/generated/rest/types/reidentified_file_output.py b/v2/skyflow/generated/rest/types/reidentified_file_output.py
similarity index 100%
rename from skyflow/generated/rest/types/reidentified_file_output.py
rename to v2/skyflow/generated/rest/types/reidentified_file_output.py
diff --git a/skyflow/generated/rest/types/reidentified_file_output_processed_file_extension.py b/v2/skyflow/generated/rest/types/reidentified_file_output_processed_file_extension.py
similarity index 100%
rename from skyflow/generated/rest/types/reidentified_file_output_processed_file_extension.py
rename to v2/skyflow/generated/rest/types/reidentified_file_output_processed_file_extension.py
diff --git a/skyflow/generated/rest/types/reidentify_file_response.py b/v2/skyflow/generated/rest/types/reidentify_file_response.py
similarity index 100%
rename from skyflow/generated/rest/types/reidentify_file_response.py
rename to v2/skyflow/generated/rest/types/reidentify_file_response.py
diff --git a/skyflow/generated/rest/types/reidentify_file_response_output_type.py b/v2/skyflow/generated/rest/types/reidentify_file_response_output_type.py
similarity index 100%
rename from skyflow/generated/rest/types/reidentify_file_response_output_type.py
rename to v2/skyflow/generated/rest/types/reidentify_file_response_output_type.py
diff --git a/skyflow/generated/rest/types/reidentify_file_response_status.py b/v2/skyflow/generated/rest/types/reidentify_file_response_status.py
similarity index 100%
rename from skyflow/generated/rest/types/reidentify_file_response_status.py
rename to v2/skyflow/generated/rest/types/reidentify_file_response_status.py
diff --git a/skyflow/generated/rest/types/request_action_type.py b/v2/skyflow/generated/rest/types/request_action_type.py
similarity index 100%
rename from skyflow/generated/rest/types/request_action_type.py
rename to v2/skyflow/generated/rest/types/request_action_type.py
diff --git a/skyflow/generated/rest/types/resource_id.py b/v2/skyflow/generated/rest/types/resource_id.py
similarity index 100%
rename from skyflow/generated/rest/types/resource_id.py
rename to v2/skyflow/generated/rest/types/resource_id.py
diff --git a/skyflow/generated/rest/types/shift_dates.py b/v2/skyflow/generated/rest/types/shift_dates.py
similarity index 100%
rename from skyflow/generated/rest/types/shift_dates.py
rename to v2/skyflow/generated/rest/types/shift_dates.py
diff --git a/skyflow/generated/rest/types/shift_dates_entity_types_item.py b/v2/skyflow/generated/rest/types/shift_dates_entity_types_item.py
similarity index 100%
rename from skyflow/generated/rest/types/shift_dates_entity_types_item.py
rename to v2/skyflow/generated/rest/types/shift_dates_entity_types_item.py
diff --git a/skyflow/generated/rest/types/string_response_entities.py b/v2/skyflow/generated/rest/types/string_response_entities.py
similarity index 100%
rename from skyflow/generated/rest/types/string_response_entities.py
rename to v2/skyflow/generated/rest/types/string_response_entities.py
diff --git a/skyflow/generated/rest/types/token_type_mapping.py b/v2/skyflow/generated/rest/types/token_type_mapping.py
similarity index 100%
rename from skyflow/generated/rest/types/token_type_mapping.py
rename to v2/skyflow/generated/rest/types/token_type_mapping.py
diff --git a/skyflow/generated/rest/types/token_type_mapping_default.py b/v2/skyflow/generated/rest/types/token_type_mapping_default.py
similarity index 100%
rename from skyflow/generated/rest/types/token_type_mapping_default.py
rename to v2/skyflow/generated/rest/types/token_type_mapping_default.py
diff --git a/skyflow/generated/rest/types/token_type_mapping_entity_only_item.py b/v2/skyflow/generated/rest/types/token_type_mapping_entity_only_item.py
similarity index 100%
rename from skyflow/generated/rest/types/token_type_mapping_entity_only_item.py
rename to v2/skyflow/generated/rest/types/token_type_mapping_entity_only_item.py
diff --git a/skyflow/generated/rest/types/token_type_mapping_entity_unq_counter_item.py b/v2/skyflow/generated/rest/types/token_type_mapping_entity_unq_counter_item.py
similarity index 100%
rename from skyflow/generated/rest/types/token_type_mapping_entity_unq_counter_item.py
rename to v2/skyflow/generated/rest/types/token_type_mapping_entity_unq_counter_item.py
diff --git a/skyflow/generated/rest/types/token_type_mapping_vault_token_item.py b/v2/skyflow/generated/rest/types/token_type_mapping_vault_token_item.py
similarity index 100%
rename from skyflow/generated/rest/types/token_type_mapping_vault_token_item.py
rename to v2/skyflow/generated/rest/types/token_type_mapping_vault_token_item.py
diff --git a/skyflow/generated/rest/types/transformations.py b/v2/skyflow/generated/rest/types/transformations.py
similarity index 100%
rename from skyflow/generated/rest/types/transformations.py
rename to v2/skyflow/generated/rest/types/transformations.py
diff --git a/skyflow/generated/rest/types/upload_file_v_2_response.py b/v2/skyflow/generated/rest/types/upload_file_v_2_response.py
similarity index 100%
rename from skyflow/generated/rest/types/upload_file_v_2_response.py
rename to v2/skyflow/generated/rest/types/upload_file_v_2_response.py
diff --git a/skyflow/generated/rest/types/uuid_.py b/v2/skyflow/generated/rest/types/uuid_.py
similarity index 100%
rename from skyflow/generated/rest/types/uuid_.py
rename to v2/skyflow/generated/rest/types/uuid_.py
diff --git a/skyflow/generated/rest/types/v_1_audit_after_options.py b/v2/skyflow/generated/rest/types/v_1_audit_after_options.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_audit_after_options.py
rename to v2/skyflow/generated/rest/types/v_1_audit_after_options.py
diff --git a/skyflow/generated/rest/types/v_1_audit_event_response.py b/v2/skyflow/generated/rest/types/v_1_audit_event_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_audit_event_response.py
rename to v2/skyflow/generated/rest/types/v_1_audit_event_response.py
diff --git a/skyflow/generated/rest/types/v_1_audit_response.py b/v2/skyflow/generated/rest/types/v_1_audit_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_audit_response.py
rename to v2/skyflow/generated/rest/types/v_1_audit_response.py
diff --git a/skyflow/generated/rest/types/v_1_audit_response_event.py b/v2/skyflow/generated/rest/types/v_1_audit_response_event.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_audit_response_event.py
rename to v2/skyflow/generated/rest/types/v_1_audit_response_event.py
diff --git a/skyflow/generated/rest/types/v_1_audit_response_event_request.py b/v2/skyflow/generated/rest/types/v_1_audit_response_event_request.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_audit_response_event_request.py
rename to v2/skyflow/generated/rest/types/v_1_audit_response_event_request.py
diff --git a/skyflow/generated/rest/types/v_1_batch_operation_response.py b/v2/skyflow/generated/rest/types/v_1_batch_operation_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_batch_operation_response.py
rename to v2/skyflow/generated/rest/types/v_1_batch_operation_response.py
diff --git a/skyflow/generated/rest/types/v_1_batch_record.py b/v2/skyflow/generated/rest/types/v_1_batch_record.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_batch_record.py
rename to v2/skyflow/generated/rest/types/v_1_batch_record.py
diff --git a/skyflow/generated/rest/types/v_1_bin_list_response.py b/v2/skyflow/generated/rest/types/v_1_bin_list_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_bin_list_response.py
rename to v2/skyflow/generated/rest/types/v_1_bin_list_response.py
diff --git a/skyflow/generated/rest/types/v_1_bulk_delete_record_response.py b/v2/skyflow/generated/rest/types/v_1_bulk_delete_record_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_bulk_delete_record_response.py
rename to v2/skyflow/generated/rest/types/v_1_bulk_delete_record_response.py
diff --git a/skyflow/generated/rest/types/v_1_bulk_get_record_response.py b/v2/skyflow/generated/rest/types/v_1_bulk_get_record_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_bulk_get_record_response.py
rename to v2/skyflow/generated/rest/types/v_1_bulk_get_record_response.py
diff --git a/skyflow/generated/rest/types/v_1_byot.py b/v2/skyflow/generated/rest/types/v_1_byot.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_byot.py
rename to v2/skyflow/generated/rest/types/v_1_byot.py
diff --git a/skyflow/generated/rest/types/v_1_card.py b/v2/skyflow/generated/rest/types/v_1_card.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_card.py
rename to v2/skyflow/generated/rest/types/v_1_card.py
diff --git a/skyflow/generated/rest/types/v_1_delete_file_response.py b/v2/skyflow/generated/rest/types/v_1_delete_file_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_delete_file_response.py
rename to v2/skyflow/generated/rest/types/v_1_delete_file_response.py
diff --git a/skyflow/generated/rest/types/v_1_delete_record_response.py b/v2/skyflow/generated/rest/types/v_1_delete_record_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_delete_record_response.py
rename to v2/skyflow/generated/rest/types/v_1_delete_record_response.py
diff --git a/skyflow/generated/rest/types/v_1_detokenize_record_request.py b/v2/skyflow/generated/rest/types/v_1_detokenize_record_request.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_detokenize_record_request.py
rename to v2/skyflow/generated/rest/types/v_1_detokenize_record_request.py
diff --git a/skyflow/generated/rest/types/v_1_detokenize_record_response.py b/v2/skyflow/generated/rest/types/v_1_detokenize_record_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_detokenize_record_response.py
rename to v2/skyflow/generated/rest/types/v_1_detokenize_record_response.py
diff --git a/skyflow/generated/rest/types/v_1_detokenize_response.py b/v2/skyflow/generated/rest/types/v_1_detokenize_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_detokenize_response.py
rename to v2/skyflow/generated/rest/types/v_1_detokenize_response.py
diff --git a/skyflow/generated/rest/types/v_1_field_records.py b/v2/skyflow/generated/rest/types/v_1_field_records.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_field_records.py
rename to v2/skyflow/generated/rest/types/v_1_field_records.py
diff --git a/skyflow/generated/rest/types/v_1_file_av_scan_status.py b/v2/skyflow/generated/rest/types/v_1_file_av_scan_status.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_file_av_scan_status.py
rename to v2/skyflow/generated/rest/types/v_1_file_av_scan_status.py
diff --git a/v2/skyflow/generated/rest/types/v_1_get_auth_token_response.py b/v2/skyflow/generated/rest/types/v_1_get_auth_token_response.py
new file mode 100644
index 00000000..c4db65a0
--- /dev/null
+++ b/v2/skyflow/generated/rest/types/v_1_get_auth_token_response.py
@@ -0,0 +1,33 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class V1GetAuthTokenResponse(UniversalBaseModel):
+ access_token: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="accessToken")] = (
+ pydantic.Field(default=None)
+ )
+ """
+ AccessToken.
+ """
+
+ token_type: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="tokenType")] = pydantic.Field(
+ default=None
+ )
+ """
+ TokenType : Bearer.
+ """
+
+ if IS_PYDANTIC_V2:
+ model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True) # type: ignore # Pydantic v2
+ else:
+
+ class Config:
+ frozen = True
+ smart_union = True
+ extra = pydantic.Extra.allow
diff --git a/skyflow/generated/rest/types/v_1_get_file_scan_status_response.py b/v2/skyflow/generated/rest/types/v_1_get_file_scan_status_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_get_file_scan_status_response.py
rename to v2/skyflow/generated/rest/types/v_1_get_file_scan_status_response.py
diff --git a/skyflow/generated/rest/types/v_1_get_query_response.py b/v2/skyflow/generated/rest/types/v_1_get_query_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_get_query_response.py
rename to v2/skyflow/generated/rest/types/v_1_get_query_response.py
diff --git a/skyflow/generated/rest/types/v_1_insert_record_response.py b/v2/skyflow/generated/rest/types/v_1_insert_record_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_insert_record_response.py
rename to v2/skyflow/generated/rest/types/v_1_insert_record_response.py
diff --git a/skyflow/generated/rest/types/v_1_member_type.py b/v2/skyflow/generated/rest/types/v_1_member_type.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_member_type.py
rename to v2/skyflow/generated/rest/types/v_1_member_type.py
diff --git a/skyflow/generated/rest/types/v_1_record_meta_properties.py b/v2/skyflow/generated/rest/types/v_1_record_meta_properties.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_record_meta_properties.py
rename to v2/skyflow/generated/rest/types/v_1_record_meta_properties.py
diff --git a/skyflow/generated/rest/types/v_1_tokenize_record_request.py b/v2/skyflow/generated/rest/types/v_1_tokenize_record_request.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_tokenize_record_request.py
rename to v2/skyflow/generated/rest/types/v_1_tokenize_record_request.py
diff --git a/skyflow/generated/rest/types/v_1_tokenize_record_response.py b/v2/skyflow/generated/rest/types/v_1_tokenize_record_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_tokenize_record_response.py
rename to v2/skyflow/generated/rest/types/v_1_tokenize_record_response.py
diff --git a/skyflow/generated/rest/types/v_1_tokenize_response.py b/v2/skyflow/generated/rest/types/v_1_tokenize_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_tokenize_response.py
rename to v2/skyflow/generated/rest/types/v_1_tokenize_response.py
diff --git a/skyflow/generated/rest/types/v_1_update_record_response.py b/v2/skyflow/generated/rest/types/v_1_update_record_response.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_update_record_response.py
rename to v2/skyflow/generated/rest/types/v_1_update_record_response.py
diff --git a/skyflow/generated/rest/types/v_1_vault_field_mapping.py b/v2/skyflow/generated/rest/types/v_1_vault_field_mapping.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_vault_field_mapping.py
rename to v2/skyflow/generated/rest/types/v_1_vault_field_mapping.py
diff --git a/skyflow/generated/rest/types/v_1_vault_schema_config.py b/v2/skyflow/generated/rest/types/v_1_vault_schema_config.py
similarity index 100%
rename from skyflow/generated/rest/types/v_1_vault_schema_config.py
rename to v2/skyflow/generated/rest/types/v_1_vault_schema_config.py
diff --git a/skyflow/generated/rest/types/word_character_count.py b/v2/skyflow/generated/rest/types/word_character_count.py
similarity index 100%
rename from skyflow/generated/rest/types/word_character_count.py
rename to v2/skyflow/generated/rest/types/word_character_count.py
diff --git a/skyflow/generated/rest/version.py b/v2/skyflow/generated/rest/version.py
similarity index 100%
rename from skyflow/generated/rest/version.py
rename to v2/skyflow/generated/rest/version.py
diff --git a/v2/skyflow/py.typed b/v2/skyflow/py.typed
new file mode 100644
index 00000000..e69de29b
diff --git a/skyflow/service_account/__init__.py b/v2/skyflow/service_account/__init__.py
similarity index 100%
rename from skyflow/service_account/__init__.py
rename to v2/skyflow/service_account/__init__.py
diff --git a/skyflow/service_account/_utils.py b/v2/skyflow/service_account/_utils.py
similarity index 100%
rename from skyflow/service_account/_utils.py
rename to v2/skyflow/service_account/_utils.py
diff --git a/v2/skyflow/service_account/client/__init__.py b/v2/skyflow/service_account/client/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/skyflow/service_account/client/auth_client.py b/v2/skyflow/service_account/client/auth_client.py
similarity index 100%
rename from skyflow/service_account/client/auth_client.py
rename to v2/skyflow/service_account/client/auth_client.py
diff --git a/skyflow/utils/__init__.py b/v2/skyflow/utils/__init__.py
similarity index 100%
rename from skyflow/utils/__init__.py
rename to v2/skyflow/utils/__init__.py
diff --git a/v2/skyflow/utils/_helpers.py b/v2/skyflow/utils/_helpers.py
new file mode 100644
index 00000000..12ff1257
--- /dev/null
+++ b/v2/skyflow/utils/_helpers.py
@@ -0,0 +1,18 @@
+from urllib.parse import urlparse
+
+def get_base_url(url):
+ parsed_url = urlparse(url)
+ base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
+ return base_url
+
+def format_scope(scopes):
+ if not scopes:
+ return None
+ return " ".join([f"role:{scope}" for scope in scopes])
+
+def is_valid_url(url):
+ try:
+ result = urlparse(url)
+ return all([result.scheme == "https", result.netloc])
+ except Exception:
+ return False
\ No newline at end of file
diff --git a/skyflow/utils/_skyflow_messages.py b/v2/skyflow/utils/_skyflow_messages.py
similarity index 99%
rename from skyflow/utils/_skyflow_messages.py
rename to v2/skyflow/utils/_skyflow_messages.py
index 232bd8b0..9149e1d4 100644
--- a/skyflow/utils/_skyflow_messages.py
+++ b/v2/skyflow/utils/_skyflow_messages.py
@@ -91,6 +91,9 @@ class Error(Enum):
INVALID_TABLE_NAME_IN_INSERT = f"{error_prefix} Validation error. Invalid table name in insert request. Specify a valid table name."
INVALID_TYPE_OF_DATA_IN_INSERT = f"{error_prefix} Validation error. Invalid type of data in insert request. Specify data as a object array."
EMPTY_DATA_IN_INSERT = f"{error_prefix} Validation error. Data array cannot be empty. Specify data in insert request."
+ INVALID_RECORD_DATA_IN_INSERT = f"{error_prefix} Validation error. Each record's field values must be a non-empty dict."
+ EMPTY_KEY_IN_INSERT_DATA = f"{error_prefix} Validation error. A record must not contain a null or empty key."
+ EMPTY_VALUE_IN_INSERT_DATA = f"{error_prefix} Validation error. A record must not contain a null or empty value."
INVALID_UPSERT_OPTIONS_TYPE = f"{error_prefix} Validation error. Invalid 'upsert' value in options. Specify 'upsert' as a non-empty string containing the column name."
INVALID_HOMOGENEOUS_TYPE = f"{error_prefix} Validation error. Invalid type of homogeneous. Specify homogeneous as a string."
INVALID_TOKEN_MODE_TYPE = f"{error_prefix} Validation error. Invalid type of token mode. Specify token mode as a TokenMode enum."
diff --git a/skyflow/utils/_utils.py b/v2/skyflow/utils/_utils.py
similarity index 100%
rename from skyflow/utils/_utils.py
rename to v2/skyflow/utils/_utils.py
diff --git a/skyflow/utils/_version.py b/v2/skyflow/utils/_version.py
similarity index 100%
rename from skyflow/utils/_version.py
rename to v2/skyflow/utils/_version.py
diff --git a/v2/skyflow/utils/constants.py b/v2/skyflow/utils/constants.py
new file mode 100644
index 00000000..05d28380
--- /dev/null
+++ b/v2/skyflow/utils/constants.py
@@ -0,0 +1,291 @@
+OPTIONAL_TOKEN='token'
+PROTOCOL='https'
+SKY_META_DATA_HEADER='sky-metadata'
+CTX_KEY_REGEX=r'^[a-zA-Z0-9_]+$'
+
+class SKYFLOW:
+ SKYFLOW_ID = 'skyflowId'
+ X_SKYFLOW_AUTHORIZATION = 'x-skyflow-authorization'
+
+
+class HttpHeader:
+ CONTENT_TYPE = 'Content-Type'
+ CONTENT_TYPE_LOWERCASE = 'content-type'
+ X_REQUEST_ID = 'x-request-id'
+ ERROR_FROM_CLIENT = 'error-from-client'
+ AUTHORIZATION = 'Authorization'
+ X_SKYFLOW_AUTHORIZATION_HEADER = 'X-Skyflow-Authorization'
+
+
+class HttpStatusCode:
+ OK = 200
+ BAD_REQUEST = 400
+ UNAUTHORIZED = 401
+ INTERNAL_SERVER_ERROR = 500
+
+
+class ContentType:
+ APPLICATION_JSON = 'application/json'
+ APPLICATION_X_WWW_FORM_URLENCODED = 'application/x-www-form-urlencoded'
+ TEXT_PLAIN = 'text/plain'
+
+
+class DetectStatus:
+ IN_PROGRESS = 'IN_PROGRESS'
+ SUCCESS = 'SUCCESS'
+ FAILED = 'FAILED'
+ UNKNOWN = 'UNKNOWN'
+
+class Detect:
+ WAIT_TIME = 64
+
+class FileExtension:
+ JSON = 'json'
+ MP3 = 'mp3'
+ WAV = 'wav'
+ PDF = 'pdf'
+ TXT = 'txt'
+ DOC = 'doc'
+ DOCX = 'docx'
+ JPG = 'jpg'
+ JPEG = 'jpeg'
+ PNG = 'png'
+ BMP = 'bmp'
+ TIF = 'tif'
+ TIFF = 'tiff'
+ PPT = 'ppt'
+ PPTX = 'pptx'
+ CSV = 'csv'
+ XLS = 'xls'
+ XLSX = 'xlsx'
+ XML = 'xml'
+
+
+class FileProcessing:
+ PROCESSED_PREFIX = 'processed-'
+ DEIDENTIFIED_PREFIX = 'deidentified.'
+ ENTITIES = 'entities'
+
+
+class EncodingType:
+ UTF8 = 'utf8'
+ UTF_8 = 'utf-8'
+ BASE64 = 'base64'
+ BINARY = 'binary'
+
+
+class JWT:
+ ALGORITHM_RS256 = 'RS256'
+ GRANT_TYPE_JWT_BEARER = 'urn:ietf:params:oauth:grant-type:jwt-bearer'
+ ISSUER_SDK = 'sdk'
+ SIGNED_TOKEN_PREFIX = 'signed_token_'
+ ROLE_PREFIX = 'role:'
+
+
+class ApiKey:
+ SKY_PREFIX = 'sky-'
+ LENGTH = 42
+
+
+class UrlProtocol:
+ HTTPS = 'https'
+ HTTP = 'http'
+
+
+class BooleanString:
+ TRUE = 'true'
+ FALSE = 'false'
+
+
+class ResponseField:
+ STATUS = 'Status'
+ BODY = 'Body'
+ RECORDS = 'records'
+ TOKENS = 'tokens'
+ ERROR = 'error'
+ SKYFLOW_ID = 'skyflow_id'
+ REQUEST_INDEX = 'request_index'
+ REQUEST_ID = 'request_id'
+ HTTP_CODE = 'http_code'
+ HTTP_STATUS = 'http_status'
+ GRPC_CODE = 'grpc_code'
+ DETAILS = 'details'
+ MESSAGE = 'message'
+ ERROR_FROM_CLIENT = 'error_from_client'
+ TOKEN = 'token'
+ VALUE = 'value'
+ TYPE = 'type'
+ TOKENIZED_DATA = 'tokenized_data'
+ SIGNED_TOKEN = 'signed_token'
+ RESPONSES = 'responses'
+
+
+class CredentialField:
+ PRIVATE_KEY = 'privateKey'
+ CLIENT_ID = 'clientID'
+ KEY_ID = 'keyID'
+ TOKEN_URI = 'tokenURI'
+ TOKEN_URI_OPTION = 'token_uri'
+ CLIENT_NAME = 'clientName'
+ CREDENTIALS_STRING = 'credentials_string'
+ API_KEY = 'api_key'
+ TOKEN = 'token'
+ PATH = 'path'
+ CONTEXT = 'context'
+ ROLES = 'roles'
+
+
+class JwtField:
+ ISS = 'iss'
+ KEY = 'key'
+ AUD = 'aud'
+ SUB = 'sub'
+ EXP = 'exp'
+ CTX = 'ctx'
+ TOK = 'tok'
+ IAT = 'iat'
+
+
+class OptionField:
+ ROLE_IDS = 'role_ids'
+ DATA_TOKENS = 'data_tokens'
+ TIME_TO_LIVE = 'time_to_live'
+ ROLES = 'roles'
+ CTX = 'ctx'
+ VAULT_ID = 'vault_id'
+ CONNECTION_ID = 'connection_id'
+ CONNECTION_URL = 'connection_url'
+ VAULT_CLIENT = 'vault_client'
+ VAULT_CONTROLLER = 'vault_controller'
+ DETECT_CONTROLLER = 'detect_controller'
+ CONTROLLER = 'controller'
+ VERIFY_SIGNATURE = 'verify_signature'
+ VERIFY_AUD = 'verify_aud'
+
+
+class ConfigField:
+ CREDENTIALS = 'credentials'
+ CLUSTER_ID = 'cluster_id'
+ ENV = 'env'
+ VAULT_ID = 'vault_id'
+
+
+class RequestParameter:
+ VALUE = 'value'
+ COLUMN_GROUP = 'column_group'
+ REDACTION = 'redaction'
+ REDACTION_TYPE = 'redaction_type'
+
+
+class FileUploadField:
+ TABLE = 'table'
+ SKYFLOW_ID = 'skyflow_id'
+ COLUMN_NAME = 'column_name'
+ FILE_PATH = 'file_path'
+ BASE64 = 'base64'
+ FILE_OBJECT = 'file_object'
+ FILE_NAME = 'file_name'
+ FILE = 'file'
+ NAME = 'name'
+
+
+class DeidentifyFileRequestField:
+ ENTITIES = 'entities'
+ ALLOW_REGEX_LIST = 'allow_regex_list'
+ RESTRICT_REGEX_LIST = 'restrict_regex_list'
+ OUTPUT_PROCESSED_IMAGE = 'output_processed_image'
+ OUTPUT_OCR_TEXT = 'output_ocr_text'
+ MASKING_METHOD = 'masking_method'
+ PIXEL_DENSITY = 'pixel_density'
+ DENSITY = 'density'
+ MAX_RESOLUTION = 'max_resolution'
+ OUTPUT_PROCESSED_AUDIO = 'output_processed_audio'
+ OUTPUT_TRANSCRIPTION = 'output_transcription'
+ BLEEP = 'bleep'
+ OUTPUT_DIRECTORY = 'output_directory'
+ WAIT_TIME = 'wait_time'
+
+
+class DeidentifyField:
+ TEXT = 'text'
+ ENTITY_TYPES = 'entity_types'
+ TOKEN_TYPE = 'token_type'
+ ALLOW_REGEX = 'allow_regex'
+ RESTRICT_REGEX = 'restrict_regex'
+ TRANSFORMATIONS = 'transformations'
+ FORMAT = 'format'
+ OUTPUT = 'output'
+ STATUS = 'status'
+ RUN_ID = 'run_id'
+ WORD_CHARACTER_COUNT = 'word_character_count'
+ WORD_COUNT = 'word_count'
+ CHARACTER_COUNT = 'character_count'
+ SIZE = 'size'
+ DURATION = 'duration'
+ PAGES = 'pages'
+ SLIDES = 'slides'
+ PROCESSED_FILE = 'processed_file'
+ PROCESSED_FILE_TYPE = 'processed_file_type'
+ PROCESSED_FILE_EXTENSION = 'processed_file_extension'
+ REDACTED_FILE = 'redacted_file'
+ SHIFT_DATES = 'shift_dates'
+ DEFAULT = 'default'
+ ENTITY_UNQ_COUNTER = 'entity_unq_counter'
+ ENTITY_UNIQUE_COUNTER = 'entity_unique_counter'
+ ENTITY_ONLY = 'entity_only'
+ VAULT_TOKEN = 'vault_token'
+ ENTITIES = 'entities'
+ MAX_DAYS = 'max_days'
+ MIN_DAYS = 'min_days'
+ MAX = 'max'
+ MIN = 'min'
+ FILE = 'file'
+ TYPE = 'type'
+ EXTENSION = 'extension'
+ IN_PROGRESS = 'IN_PROGRESS'
+ REQUEST_OPTIONS = 'request_options'
+ BLEEP_GAIN = 'bleep_gain'
+ BLEEP_FREQUENCY = 'bleep_frequency'
+ BLEEP_START_PADDING = 'bleep_start_padding'
+ BLEEP_STOP_PADDING = 'bleep_stop_padding'
+ DENSITY = 'density'
+ TOKEN_FORMAT = 'token_format'
+ PROCESSED_FILE_RESPONSE_KEY = 'processedFile'
+ PROCESSED_FILE_TYPE_RESPONSE_KEY = 'processedFileType'
+ PROCESSED_FILE_EXTENSION_RESPONSE_KEY = 'processedFileExtension'
+
+
+class RequestOperation:
+ INSERT = 'INSERT'
+ DELETE = 'DELETE'
+ GET = 'GET'
+ UPDATE = 'UPDATE'
+ QUERY = 'QUERY'
+ TOKENIZE = 'TOKENIZE'
+ DETOKENIZE = 'DETOKENIZE'
+ FILE_UPLOAD = 'FILE_UPLOAD'
+
+
+class ConfigType:
+ VAULT = 'vault'
+ CONNECTION = 'connection'
+
+
+class SqlCommand:
+ SELECT = 'SELECT'
+
+
+class SdkPrefix:
+ SKYFLOW_PYTHON = 'skyflow-python@'
+ PYTHON_RUNTIME = 'Python '
+
+
+class SdkMetricsKey:
+ SDK_NAME_VERSION = 'sdk_name_version'
+ SDK_CLIENT_DEVICE_MODEL = 'sdk_client_device_model'
+ SDK_CLIENT_OS_DETAILS = 'sdk_client_os_details'
+ SDK_RUNTIME_DETAILS = 'sdk_runtime_details'
+
+
+class ErrorDefaults:
+ UNKNOWN_REQUEST_ID = 'unknown-request-id'
diff --git a/v2/skyflow/utils/enums/__init__.py b/v2/skyflow/utils/enums/__init__.py
new file mode 100644
index 00000000..af293ce2
--- /dev/null
+++ b/v2/skyflow/utils/enums/__init__.py
@@ -0,0 +1,12 @@
+from .env import Env, EnvUrls
+from .log_level import LogLevel
+from .content_types import ContentType
+from .detect_entities import DetectEntities
+from .token_mode import TokenMode
+from .token_type import TokenType
+from .request_method import RequestMethod
+from .redaction_type import RedactionType
+from .detect_entities import DetectEntities
+from .detect_output_transcriptions import DetectOutputTranscriptions
+from .masking_method import MaskingMethod
+from .token_type import TokenType
\ No newline at end of file
diff --git a/v2/skyflow/utils/enums/content_types.py b/v2/skyflow/utils/enums/content_types.py
new file mode 100644
index 00000000..f2db5b92
--- /dev/null
+++ b/v2/skyflow/utils/enums/content_types.py
@@ -0,0 +1,9 @@
+from enum import Enum
+
+class ContentType(Enum):
+ JSON = 'application/json'
+ PLAINTEXT = 'text/plain'
+ XML = 'text/xml'
+ URLENCODED = 'application/x-www-form-urlencoded'
+ FORMDATA = 'multipart/form-data'
+ HTML = 'text/html'
\ No newline at end of file
diff --git a/v2/skyflow/utils/enums/detect_entities.py b/v2/skyflow/utils/enums/detect_entities.py
new file mode 100644
index 00000000..91d5e1a2
--- /dev/null
+++ b/v2/skyflow/utils/enums/detect_entities.py
@@ -0,0 +1,73 @@
+from enum import Enum
+
+class DetectEntities(Enum):
+ ACCOUNT_NUMBER = "account_number"
+ AGE = "age"
+ ALL = "all"
+ BANK_ACCOUNT = "bank_account"
+ BLOOD_TYPE = "blood_type"
+ CONDITION = "condition"
+ CORPORATE_ACTION = "corporate_action"
+ CREDIT_CARD = "credit_card"
+ CREDIT_CARD_EXPIRATION = "credit_card_expiration"
+ CVV = "cvv"
+ DATE = "date"
+ DAY = "day"
+ DATE_INTERVAL = "date_interval"
+ DOB = "dob"
+ DOSE = "dose"
+ DRIVER_LICENSE = "driver_license"
+ DRUG = "drug"
+ DURATION = "duration"
+ EFFECT = "effect"
+ EMAIL_ADDRESS = "email_address"
+ EVENT = "event"
+ FILENAME = "filename"
+ FINANCIAL_METRIC = "financial_metric"
+ GENDER = "gender"
+ HEALTHCARE_NUMBER = "healthcare_number"
+ INJURY = "injury"
+ IP_ADDRESS = "ip_address"
+ LANGUAGE = "language"
+ LOCATION = "location"
+ LOCATION_ADDRESS = "location_address"
+ LOCATION_ADDRESS_STREET = "location_address_street"
+ LOCATION_CITY = "location_city"
+ LOCATION_COORDINATE = "location_coordinate"
+ LOCATION_COUNTRY = "location_country"
+ LOCATION_STATE = "location_state"
+ LOCATION_ZIP = "location_zip"
+ MARITAL_STATUS = "marital_status"
+ MEDICAL_CODE = "medical_code"
+ MEDICAL_PROCESS = "medical_process"
+ MONEY = "money"
+ MONTH = "month"
+ NAME = "name"
+ NAME_FAMILY = "name_family"
+ NAME_GIVEN = "name_given"
+ NAME_MEDICAL_PROFESSIONAL = "name_medical_professional"
+ NUMERICAL_PII = "numerical_pii"
+ OCCUPATION = "occupation"
+ ORGANIZATION = "organization"
+ ORGANIZATION_ID = "organization_id"
+ ORGANIZATION_MEDICAL_FACILITY = "organization_medical_facility"
+ ORIGIN = "origin"
+ PASSPORT_NUMBER = "passport_number"
+ PASSWORD = "password"
+ PHONE_NUMBER = "phone_number"
+ PROJECT = "project"
+ PHYSICAL_ATTRIBUTE = "physical_attribute"
+ POLITICAL_AFFILIATION = "political_affiliation"
+ PRODUCT = "product"
+ RELIGION = "religion"
+ ROUTING_NUMBER = "routing_number"
+ SEXUALITY = "sexuality"
+ SSN = "ssn"
+ STATISTICS = "statistics"
+ TIME = "time"
+ TREND = "trend"
+ URL = "url"
+ USERNAME = "username"
+ VEHICLE_ID = "vehicle_id"
+ YEAR = "year"
+ ZODIAC_SIGN = "zodiac_sign"
\ No newline at end of file
diff --git a/v2/skyflow/utils/enums/detect_output_transcriptions.py b/v2/skyflow/utils/enums/detect_output_transcriptions.py
new file mode 100644
index 00000000..a398a3d8
--- /dev/null
+++ b/v2/skyflow/utils/enums/detect_output_transcriptions.py
@@ -0,0 +1,8 @@
+from enum import Enum
+
+class DetectOutputTranscriptions(Enum):
+ DIARIZED_TRANSCRIPTION = "diarized_transcription"
+ MEDICAL_DIARIZED_TRANSCRIPTION = "medical_diarized_transcription"
+ MEDICAL_TRANSCRIPTION = "medical_transcription"
+ TRANSCRIPTION = "transcription"
+ PLAINTEXT_TRANSCRIPTION = "plaintext_transcription"
\ No newline at end of file
diff --git a/v2/skyflow/utils/enums/env.py b/v2/skyflow/utils/enums/env.py
new file mode 100644
index 00000000..512812b8
--- /dev/null
+++ b/v2/skyflow/utils/enums/env.py
@@ -0,0 +1,16 @@
+# Re-exports common's Env/EnvUrls rather than defining a separate duplicate class.
+#
+# Why this matters (found via a live run, not caught by tests): VaultClient.initialize_client_configuration()
+# is inherited from common.vault.base_vault_client.BaseVaultClient, which calls common.utils.get_vault_url().
+# That function validates its `env` argument with `if env not in Env` against *common's* Env class. If this
+# module defined its own separate (even if identically-shaped) Env class, a value built from `skyflow.Env`
+# would never satisfy that check -- plain Enum classes never compare equal across distinct class objects,
+# even with matching member names/values. Re-exporting the same class object avoids that entirely.
+#
+# This is safe for v2 compatibility: values, names, and `.value`/`.name` behavior are unchanged --
+# `skyflow.Env.PROD` still is `Env.PROD` with the same string value. Only the class's identity is now
+# shared instead of duplicated, which is exactly what a config value needs to survive a round trip through
+# shared validation code in common/.
+from common.utils.enums import Env, EnvUrls
+
+__all__ = ["Env", "EnvUrls"]
diff --git a/v2/skyflow/utils/enums/log_level.py b/v2/skyflow/utils/enums/log_level.py
new file mode 100644
index 00000000..42aa1dfe
--- /dev/null
+++ b/v2/skyflow/utils/enums/log_level.py
@@ -0,0 +1 @@
+from common.utils.enums.log_level import LogLevel
diff --git a/v2/skyflow/utils/enums/masking_method.py b/v2/skyflow/utils/enums/masking_method.py
new file mode 100644
index 00000000..a322f35f
--- /dev/null
+++ b/v2/skyflow/utils/enums/masking_method.py
@@ -0,0 +1,5 @@
+from enum import Enum
+
+class MaskingMethod(Enum):
+ BLACKBOX= "blackbox"
+ BLUR= "blur"
\ No newline at end of file
diff --git a/v2/skyflow/utils/enums/redaction_type.py b/v2/skyflow/utils/enums/redaction_type.py
new file mode 100644
index 00000000..1780e820
--- /dev/null
+++ b/v2/skyflow/utils/enums/redaction_type.py
@@ -0,0 +1,7 @@
+from enum import Enum
+
+class RedactionType(Enum):
+ PLAIN_TEXT = 'PLAIN_TEXT'
+ MASKED = 'MASKED'
+ DEFAULT = 'DEFAULT'
+ REDACTED = 'REDACTED'
diff --git a/v2/skyflow/utils/enums/request_method.py b/v2/skyflow/utils/enums/request_method.py
new file mode 100644
index 00000000..61efef3d
--- /dev/null
+++ b/v2/skyflow/utils/enums/request_method.py
@@ -0,0 +1,8 @@
+from enum import Enum
+
+class RequestMethod(Enum):
+ GET = "GET"
+ POST = "POST"
+ PUT = "PUT"
+ DELETE = "DELETE"
+ NONE = "NONE"
\ No newline at end of file
diff --git a/v2/skyflow/utils/enums/token_mode.py b/v2/skyflow/utils/enums/token_mode.py
new file mode 100644
index 00000000..a073b125
--- /dev/null
+++ b/v2/skyflow/utils/enums/token_mode.py
@@ -0,0 +1,6 @@
+from enum import Enum
+
+class TokenMode(Enum):
+ DISABLE = "DISABLE"
+ ENABLE = "ENABLE"
+ ENABLE_STRICT = "ENABLE_STRICT"
\ No newline at end of file
diff --git a/v2/skyflow/utils/enums/token_type.py b/v2/skyflow/utils/enums/token_type.py
new file mode 100644
index 00000000..9e9e5fcf
--- /dev/null
+++ b/v2/skyflow/utils/enums/token_type.py
@@ -0,0 +1,6 @@
+from enum import Enum
+
+class TokenType(Enum):
+ VAULT_TOKEN = "vault_token"
+ ENTITY_UNIQUE_COUNTER = "entity_unq_counter"
+ ENTITY_ONLY = "entity_only"
diff --git a/v2/skyflow/utils/logger/__init__.py b/v2/skyflow/utils/logger/__init__.py
new file mode 100644
index 00000000..bce55608
--- /dev/null
+++ b/v2/skyflow/utils/logger/__init__.py
@@ -0,0 +1,2 @@
+from ._logger import Logger
+from ._log_helpers import log_error, log_info, log_warn, log_error_log, set_active_log_level
\ No newline at end of file
diff --git a/v2/skyflow/utils/logger/_log_helpers.py b/v2/skyflow/utils/logger/_log_helpers.py
new file mode 100644
index 00000000..4fdff077
--- /dev/null
+++ b/v2/skyflow/utils/logger/_log_helpers.py
@@ -0,0 +1 @@
+from common.utils.logger._log_helpers import log_error, log_info, log_warn, log_error_log, set_active_log_level
diff --git a/v2/skyflow/utils/logger/_logger.py b/v2/skyflow/utils/logger/_logger.py
new file mode 100644
index 00000000..423df426
--- /dev/null
+++ b/v2/skyflow/utils/logger/_logger.py
@@ -0,0 +1 @@
+from common.utils.logger._logger import Logger
diff --git a/skyflow/utils/validations/__init__.py b/v2/skyflow/utils/validations/__init__.py
similarity index 100%
rename from skyflow/utils/validations/__init__.py
rename to v2/skyflow/utils/validations/__init__.py
diff --git a/skyflow/utils/validations/_validations.py b/v2/skyflow/utils/validations/_validations.py
similarity index 82%
rename from skyflow/utils/validations/_validations.py
rename to v2/skyflow/utils/validations/_validations.py
index 42abe188..954b897b 100644
--- a/skyflow/utils/validations/_validations.py
+++ b/v2/skyflow/utils/validations/_validations.py
@@ -1,8 +1,14 @@
import base64
import json
import os
+from common.utils.validations import (
+ validate_credentials as _common_validate_credentials,
+ validate_log_level as _common_validate_log_level,
+ validate_vault_config as _common_validate_vault_config,
+ validate_update_vault_config as _common_validate_update_vault_config,
+)
from skyflow.service_account import is_expired
-from skyflow.utils.enums import LogLevel, Env, RedactionType, TokenMode, DetectEntities, DetectOutputTranscriptions, \
+from skyflow.utils.enums import RedactionType, TokenMode, DetectEntities, DetectOutputTranscriptions, \
MaskingMethod
from skyflow.error import SkyflowError
from skyflow.utils import SkyflowMessages
@@ -17,12 +23,6 @@
from skyflow.vault.detect._file_input import FileInput
from skyflow.utils._helpers import is_valid_url
-valid_vault_config_keys = [
- ConfigField.VAULT_ID,
- ConfigField.CLUSTER_ID,
- ConfigField.CREDENTIALS,
- ConfigField.ENV
-]
valid_connection_config_keys = [
OptionField.CONNECTION_ID,
OptionField.CONNECTION_URL,
@@ -82,101 +82,14 @@ def validate_api_key(api_key: str, logger = None) -> bool:
return True
def validate_credentials(logger, credentials, config_id_type=None, config_id=None):
- key_present = [k for k in [CredentialField.PATH, CredentialField.TOKEN, CredentialField.CREDENTIALS_STRING, CredentialField.API_KEY] if credentials.get(k)]
-
- if len(key_present) == 0:
- error_message = (
- SkyflowMessages.Error.INVALID_CREDENTIALS_IN_CONFIG.value.format(config_id_type, config_id)
- if config_id_type and config_id else
- SkyflowMessages.Error.INVALID_CREDENTIALS.value
- )
- log_error_log(error_message, logger)
- raise SkyflowError(error_message, invalid_input_error_code)
- elif len(key_present) > 1:
- error_message = (
- SkyflowMessages.Error.MULTIPLE_CREDENTIALS_PASSED_IN_CONFIG.value.format(config_id_type, config_id)
- if config_id_type and config_id else
- SkyflowMessages.Error.MULTIPLE_CREDENTIALS_PASSED.value
- )
- log_error_log(error_message, logger)
- raise SkyflowError(error_message, invalid_input_error_code)
-
- if CredentialField.ROLES in credentials:
- validate_required_field(
- logger, credentials, CredentialField.ROLES, list,
- SkyflowMessages.Error.INVALID_ROLES_KEY_TYPE_IN_CONFIG.value.format(config_id_type, config_id)
- if config_id_type and config_id else SkyflowMessages.Error.INVALID_ROLES_KEY_TYPE.value,
- SkyflowMessages.Error.EMPTY_ROLES_IN_CONFIG.value.format(config_id_type, config_id)
- if config_id_type and config_id else SkyflowMessages.Error.EMPTY_ROLES.value
- )
-
- if CredentialField.CONTEXT in credentials:
- validate_required_field(
- logger, credentials, CredentialField.CONTEXT, str,
- SkyflowMessages.Error.EMPTY_CONTEXT_IN_CONFIG.value.format(config_id_type, config_id)
- if config_id_type and config_id else SkyflowMessages.Error.EMPTY_CONTEXT.value,
- SkyflowMessages.Error.INVALID_CONTEXT_IN_CONFIG.value.format(config_id_type, config_id)
- if config_id_type and config_id else SkyflowMessages.Error.INVALID_CONTEXT.value
- )
-
- if CredentialField.CREDENTIALS_STRING in credentials:
- validate_required_field(
- logger, credentials, CredentialField.CREDENTIALS_STRING, str,
- SkyflowMessages.Error.EMPTY_CREDENTIALS_STRING_IN_CONFIG.value.format(config_id_type, config_id)
- if config_id_type and config_id else SkyflowMessages.Error.EMPTY_CREDENTIALS_STRING.value,
- SkyflowMessages.Error.INVALID_CREDENTIALS_STRING_IN_CONFIG.value.format(config_id_type, config_id)
- if config_id_type and config_id else SkyflowMessages.Error.INVALID_CREDENTIALS_STRING.value
- )
- elif CredentialField.PATH in credentials:
- validate_required_field(
- logger, credentials, CredentialField.PATH, str,
- SkyflowMessages.Error.EMPTY_CREDENTIAL_FILE_PATH_IN_CONFIG.value.format(config_id_type, config_id)
- if config_id_type and config_id else SkyflowMessages.Error.EMPTY_CREDENTIAL_FILE_PATH.value,
- SkyflowMessages.Error.INVALID_CREDENTIAL_FILE_PATH_IN_CONFIG.value.format(config_id_type, config_id)
- if config_id_type and config_id else SkyflowMessages.Error.INVALID_CREDENTIAL_FILE_PATH.value
- )
- elif CredentialField.TOKEN in credentials:
- validate_required_field(
- logger, credentials, CredentialField.TOKEN, str,
- SkyflowMessages.Error.EMPTY_CREDENTIALS_TOKEN.value.format(config_id_type, config_id)
- if config_id_type and config_id else SkyflowMessages.Error.EMPTY_CREDENTIALS_TOKEN.value,
- SkyflowMessages.Error.INVALID_CREDENTIALS_TOKEN.value.format(config_id_type, config_id)
- if config_id_type and config_id else SkyflowMessages.Error.INVALID_CREDENTIALS_TOKEN.value
- )
- if is_expired(credentials.get(CredentialField.TOKEN), logger):
- log_error_log(SkyflowMessages.ErrorLogs.INVALID_BEARER_TOKEN.value, logger)
- raise SkyflowError(
- SkyflowMessages.Error.EXPIRED_BEARER_TOKEN.value
- if config_id_type and config_id else SkyflowMessages.Error.EXPIRED_BEARER_TOKEN.value,
- invalid_input_error_code
- )
- elif CredentialField.API_KEY in credentials:
- validate_required_field(
- logger, credentials, CredentialField.API_KEY, str,
- SkyflowMessages.Error.EMPTY_API_KEY.value.format(config_id_type, config_id)
- if config_id_type and config_id else SkyflowMessages.Error.EMPTY_API_KEY.value,
- SkyflowMessages.Error.INVALID_API_KEY.value.format(config_id_type, config_id)
- if config_id_type and config_id else SkyflowMessages.Error.INVALID_API_KEY.value
- )
- if not validate_api_key(credentials.get(CredentialField.API_KEY), logger):
- raise SkyflowError(SkyflowMessages.Error.INVALID_API_KEY.value.format(config_id_type, config_id)
- if config_id_type and config_id else SkyflowMessages.Error.INVALID_API_KEY.value,
- invalid_input_error_code)
-
- if CredentialField.TOKEN_URI_OPTION in credentials:
- token_uri = credentials.get(CredentialField.TOKEN_URI_OPTION)
- if (
- token_uri is None
- or not isinstance(token_uri, str)
- or not is_valid_url(token_uri)
- ):
- log_error_log(SkyflowMessages.ErrorLogs.INVALID_TOKEN_URI.value, logger)
- raise SkyflowError(SkyflowMessages.Error.INVALID_TOKEN_URI.value, invalid_input_error_code)
+ """Delegates to common.utils.validations.validate_credentials -- identical logic, v2's own
+ SkyflowMessages passed through so raised error text keeps showing v2's SDK version. Kept
+ under this name/signature since validate_connection_config/validate_update_connection_config
+ (v2-only, not shared) still call it directly."""
+ return _common_validate_credentials(logger, credentials, config_id_type, config_id, messages=SkyflowMessages)
def validate_log_level(logger, log_level):
- if not isinstance(log_level, LogLevel):
- log_error_log(SkyflowMessages.ErrorLogs.INVALID_LOG_LEVEL.value, logger)
- raise SkyflowError(SkyflowMessages.Error.INVALID_LOG_LEVEL.value, invalid_input_error_code)
+ return _common_validate_log_level(logger, log_level, messages=SkyflowMessages)
def validate_keys(logger, config, config_keys):
for key in config.keys():
@@ -185,62 +98,12 @@ def validate_keys(logger, config, config_keys):
raise SkyflowError(SkyflowMessages.Error.INVALID_KEY.value.format(key), invalid_input_error_code)
def validate_vault_config(logger, config):
- log_info(SkyflowMessages.Info.VALIDATING_VAULT_CONFIG.value, logger)
- validate_keys(logger, config, valid_vault_config_keys)
-
- # Validate vault_id (string, not empty)
- validate_required_field(
- logger, config, ConfigField.VAULT_ID, str,
- SkyflowMessages.Error.EMPTY_VAULT_ID.value,
- SkyflowMessages.Error.INVALID_VAULT_ID.value
- )
- vault_id = config.get(ConfigField.VAULT_ID)
- # Validate cluster_id (string, not empty)
- validate_required_field(
- logger, config, ConfigField.CLUSTER_ID, str,
- SkyflowMessages.Error.EMPTY_CLUSTER_ID.value.format(vault_id),
- SkyflowMessages.Error.INVALID_CLUSTER_ID.value.format(vault_id)
- )
-
- # Validate credentials (dict, not empty)
- if ConfigField.CREDENTIALS in config and not config.get(ConfigField.CREDENTIALS):
- raise SkyflowError(SkyflowMessages.Error.EMPTY_CREDENTIALS.value.format(ConfigType.VAULT, vault_id), invalid_input_error_code)
-
- if ConfigField.CREDENTIALS in config and config.get(ConfigField.CREDENTIALS):
- validate_credentials(logger, config.get(ConfigField.CREDENTIALS), ConfigType.VAULT, vault_id)
-
- # Validate env (optional, should be one of LogLevel values)
- if ConfigField.ENV in config and config.get(ConfigField.ENV) not in Env:
- log_error_log(SkyflowMessages.ErrorLogs.ENV_IS_REQUIRED.value, logger)
- raise SkyflowError(SkyflowMessages.Error.INVALID_ENV.value.format(vault_id), invalid_input_error_code)
-
- return True
+ """Delegates to common.utils.validations.validate_vault_config -- identical logic to
+ flowvault's own validate_vault_config, confirmed field-for-field."""
+ return _common_validate_vault_config(logger, config, messages=SkyflowMessages)
def validate_update_vault_config(logger, config):
-
- validate_keys(logger, config, valid_vault_config_keys)
-
- # Validate vault_id (string, not empty)
- validate_required_field(
- logger, config, ConfigField.VAULT_ID, str,
- SkyflowMessages.Error.EMPTY_VAULT_ID.value,
- SkyflowMessages.Error.INVALID_VAULT_ID.value
- )
-
- vault_id = config.get(ConfigField.VAULT_ID)
-
- if ConfigField.CLUSTER_ID in config and not config.get(ConfigField.CLUSTER_ID):
- raise SkyflowError(SkyflowMessages.Error.INVALID_CLUSTER_ID.value.format(vault_id), invalid_input_error_code)
-
- if ConfigField.ENV in config and config.get(ConfigField.ENV) not in Env:
- raise SkyflowError(SkyflowMessages.Error.INVALID_ENV.value.format(vault_id), invalid_input_error_code)
-
- if ConfigField.CREDENTIALS not in config:
- raise SkyflowError(SkyflowMessages.Error.EMPTY_CREDENTIALS.value.format(ConfigType.VAULT, vault_id), invalid_input_error_code)
-
- validate_credentials(logger, config.get(ConfigField.CREDENTIALS), ConfigType.VAULT, vault_id)
-
- return True
+ return _common_validate_update_vault_config(logger, config, messages=SkyflowMessages)
def validate_connection_config(logger, config):
log_info(SkyflowMessages.Info.VALIDATING_CONNECTION_CONFIG.value, logger)
@@ -469,10 +332,8 @@ def validate_insert_request(logger, request):
log_error_log(SkyflowMessages.ErrorLogs.EMPTY_VALUES.value.format(RequestOperation.INSERT), logger=logger)
raise SkyflowError(SkyflowMessages.Error.EMPTY_DATA_IN_INSERT.value, invalid_input_error_code)
- for i, item in enumerate(request.values, start=1):
- for key, value in item.items():
- if key is None or key == "":
- log_error_log(SkyflowMessages.ErrorLogs.EMPTY_OR_NULL_KEY_IN_VALUES.value.format(RequestOperation.INSERT), logger = logger)
+ # Per-record key/value emptiness is validated by the controller via the shared
+ # BaseVaultController._validate_field_values() -- not here, to avoid duplicating that logic.
if request.upsert is not None and (not isinstance(request.upsert, str) or not request.upsert.strip()):
log_error_log(SkyflowMessages.ErrorLogs.EMPTY_UPSERT.value.format(RequestOperation.INSERT), logger=logger)
diff --git a/v2/skyflow/vault/__init__.py b/v2/skyflow/vault/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/v2/skyflow/vault/client/__init__.py b/v2/skyflow/vault/client/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/v2/skyflow/vault/client/client.py b/v2/skyflow/vault/client/client.py
new file mode 100644
index 00000000..7ddb8996
--- /dev/null
+++ b/v2/skyflow/vault/client/client.py
@@ -0,0 +1,27 @@
+from common.utils import get_vault_url
+from common.vault.base_vault_client import BaseVaultClient
+from skyflow.generated.rest.client import Skyflow
+
+
+class VaultClient(BaseVaultClient):
+ def resolve_vault_url(self, cluster_id, env, vault_id, logger=None):
+ return get_vault_url(cluster_id, env, vault_id, logger=logger)
+
+ def initialize_api_client(self, vault_url, bearer_token):
+ token_provider = lambda: self._bearer_token if self._bearer_token is not None else bearer_token # noqa: E731
+ self._api_client = Skyflow(base_url=vault_url, token=token_provider)
+
+ def get_records_api(self):
+ return self._api_client.records
+
+ def get_tokens_api(self):
+ return self._api_client.tokens
+
+ def get_query_api(self):
+ return self._api_client.query
+
+ def get_detect_text_api(self):
+ return self._api_client.strings
+
+ def get_detect_file_api(self):
+ return self._api_client.files
diff --git a/skyflow/vault/connection/__init__.py b/v2/skyflow/vault/connection/__init__.py
similarity index 100%
rename from skyflow/vault/connection/__init__.py
rename to v2/skyflow/vault/connection/__init__.py
diff --git a/skyflow/vault/connection/_invoke_connection_request.py b/v2/skyflow/vault/connection/_invoke_connection_request.py
similarity index 100%
rename from skyflow/vault/connection/_invoke_connection_request.py
rename to v2/skyflow/vault/connection/_invoke_connection_request.py
diff --git a/skyflow/vault/connection/_invoke_connection_response.py b/v2/skyflow/vault/connection/_invoke_connection_response.py
similarity index 100%
rename from skyflow/vault/connection/_invoke_connection_response.py
rename to v2/skyflow/vault/connection/_invoke_connection_response.py
diff --git a/v2/skyflow/vault/controller/__init__.py b/v2/skyflow/vault/controller/__init__.py
new file mode 100644
index 00000000..2840034e
--- /dev/null
+++ b/v2/skyflow/vault/controller/__init__.py
@@ -0,0 +1,8 @@
+from ._vault import VaultController
+from ._connections import Connection
+from ._detect import Detect
+
+# Public backward-compatible name -- existing consumers do `from skyflow.vault.controller import
+# Vault`; VaultController is the canonical internal name (extends common.vault.base_vault_controller's
+# BaseVaultController), but the old public name must keep resolving to the exact same class.
+Vault = VaultController
diff --git a/skyflow/vault/controller/_audit.py b/v2/skyflow/vault/controller/_audit.py
similarity index 100%
rename from skyflow/vault/controller/_audit.py
rename to v2/skyflow/vault/controller/_audit.py
diff --git a/skyflow/vault/controller/_bin_look_up.py b/v2/skyflow/vault/controller/_bin_look_up.py
similarity index 100%
rename from skyflow/vault/controller/_bin_look_up.py
rename to v2/skyflow/vault/controller/_bin_look_up.py
diff --git a/skyflow/vault/controller/_connections.py b/v2/skyflow/vault/controller/_connections.py
similarity index 100%
rename from skyflow/vault/controller/_connections.py
rename to v2/skyflow/vault/controller/_connections.py
diff --git a/skyflow/vault/controller/_detect.py b/v2/skyflow/vault/controller/_detect.py
similarity index 100%
rename from skyflow/vault/controller/_detect.py
rename to v2/skyflow/vault/controller/_detect.py
diff --git a/skyflow/vault/controller/_vault.py b/v2/skyflow/vault/controller/_vault.py
similarity index 93%
rename from skyflow/vault/controller/_vault.py
rename to v2/skyflow/vault/controller/_vault.py
index 6c47fe3e..16e6aeac 100644
--- a/skyflow/vault/controller/_vault.py
+++ b/v2/skyflow/vault/controller/_vault.py
@@ -2,6 +2,7 @@
import json
import os
from typing import Optional
+from common.vault.base_vault_controller import BaseVaultController
from skyflow.generated.rest import V1FieldRecords, V1BatchRecord, V1TokenizeRecordRequest, \
V1DetokenizeRecordRequest
from skyflow.generated.rest.core.file import File
@@ -14,10 +15,12 @@
from skyflow.utils.logger import log_info, log_error_log
from skyflow.utils.validations import validate_insert_request, validate_delete_request, validate_query_request, \
validate_get_request, validate_update_request, validate_detokenize_request, validate_tokenize_request, validate_file_upload_request
-from skyflow.vault.data import InsertRequest, UpdateRequest, DeleteRequest, GetRequest, QueryRequest, FileUploadRequest, FileUploadResponse
-from skyflow.vault.tokens import DetokenizeRequest, TokenizeRequest
+from skyflow.vault.data import InsertRequest, InsertResponse, UpdateRequest, UpdateResponse, DeleteRequest, DeleteResponse, GetRequest, GetResponse, QueryRequest, QueryResponse, FileUploadRequest, FileUploadResponse
+from skyflow.vault.tokens import DetokenizeRequest, DetokenizeResponse, TokenizeRequest, TokenizeResponse
+
+class VaultController(BaseVaultController):
+ _skyflow_messages = SkyflowMessages
-class Vault:
def __init__(self, vault_client):
self.__vault_client = vault_client
@@ -91,9 +94,11 @@ def __get_file_for_file_upload(self, request: FileUploadRequest) -> Optional[Fil
def __get_headers(self):
return {SKY_META_DATA_HEADER: json.dumps(get_metrics())}
- def insert(self, request: InsertRequest):
+ def insert(self, request: InsertRequest) -> InsertResponse:
log_info(SkyflowMessages.Info.VALIDATE_INSERT_REQUEST.value, self.__vault_client.get_logger())
validate_insert_request(self.__vault_client.get_logger(), request)
+ for item in request.values:
+ self._validate_field_values(item)
log_info(SkyflowMessages.Info.INSERT_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
self.__initialize()
records_api = self.__vault_client.get_records_api().with_raw_response
@@ -117,7 +122,7 @@ def insert(self, request: InsertRequest):
log_error_log(SkyflowMessages.ErrorLogs.INSERT_RECORDS_REJECTED.value, self.__vault_client.get_logger())
handle_exception(e, self.__vault_client.get_logger())
- def update(self, request: UpdateRequest):
+ def update(self, request: UpdateRequest) -> UpdateResponse:
log_info(SkyflowMessages.Info.VALIDATE_UPDATE_REQUEST.value, self.__vault_client.get_logger())
validate_update_request(self.__vault_client.get_logger(), request)
log_info(SkyflowMessages.Info.UPDATE_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
@@ -144,7 +149,7 @@ def update(self, request: UpdateRequest):
log_error_log(SkyflowMessages.ErrorLogs.UPDATE_REQUEST_REJECTED.value, logger = self.__vault_client.get_logger())
handle_exception(e, self.__vault_client.get_logger())
- def delete(self, request: DeleteRequest):
+ def delete(self, request: DeleteRequest) -> DeleteResponse:
log_info(SkyflowMessages.Info.VALIDATING_DELETE_REQUEST.value, self.__vault_client.get_logger())
validate_delete_request(self.__vault_client.get_logger(), request)
log_info(SkyflowMessages.Info.DELETE_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
@@ -165,7 +170,7 @@ def delete(self, request: DeleteRequest):
log_error_log(SkyflowMessages.ErrorLogs.DELETE_REQUEST_REJECTED.value, logger = self.__vault_client.get_logger())
handle_exception(e, self.__vault_client.get_logger())
- def get(self, request: GetRequest):
+ def get(self, request: GetRequest) -> GetResponse:
log_info(SkyflowMessages.Info.VALIDATE_GET_REQUEST.value, self.__vault_client.get_logger())
validate_get_request(self.__vault_client.get_logger(), request)
log_info(SkyflowMessages.Info.GET_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
@@ -195,7 +200,7 @@ def get(self, request: GetRequest):
log_error_log(SkyflowMessages.ErrorLogs.GET_REQUEST_REJECTED.value, self.__vault_client.get_logger())
handle_exception(e, self.__vault_client.get_logger())
- def query(self, request: QueryRequest):
+ def query(self, request: QueryRequest) -> QueryResponse:
log_info(SkyflowMessages.Info.VALIDATING_QUERY_REQUEST.value, self.__vault_client.get_logger())
validate_query_request(self.__vault_client.get_logger(), request)
log_info(SkyflowMessages.Info.QUERY_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
@@ -215,7 +220,7 @@ def query(self, request: QueryRequest):
log_error_log(SkyflowMessages.ErrorLogs.QUERY_REQUEST_REJECTED.value, self.__vault_client.get_logger())
handle_exception(e, self.__vault_client.get_logger())
- def detokenize(self, request: DetokenizeRequest):
+ def detokenize(self, request: DetokenizeRequest) -> DetokenizeResponse:
log_info(SkyflowMessages.Info.VALIDATE_DETOKENIZE_REQUEST.value, self.__vault_client.get_logger())
validate_detokenize_request(self.__vault_client.get_logger(), request)
log_info(SkyflowMessages.Info.DETOKENIZE_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
@@ -243,7 +248,7 @@ def detokenize(self, request: DetokenizeRequest):
log_error_log(SkyflowMessages.ErrorLogs.DETOKENIZE_REQUEST_REJECTED.value, logger = self.__vault_client.get_logger())
handle_exception(e, self.__vault_client.get_logger())
- def tokenize(self, request: TokenizeRequest):
+ def tokenize(self, request: TokenizeRequest) -> TokenizeResponse:
log_info(SkyflowMessages.Info.VALIDATING_TOKENIZE_REQUEST.value, self.__vault_client.get_logger())
validate_tokenize_request(self.__vault_client.get_logger(), request)
log_info(SkyflowMessages.Info.TOKENIZE_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
@@ -268,7 +273,7 @@ def tokenize(self, request: TokenizeRequest):
log_error_log(SkyflowMessages.ErrorLogs.TOKENIZE_REQUEST_REJECTED.value, logger = self.__vault_client.get_logger())
handle_exception(e, self.__vault_client.get_logger())
- def upload_file(self, request: FileUploadRequest):
+ def upload_file(self, request: FileUploadRequest) -> FileUploadResponse:
log_info(SkyflowMessages.Info.FILE_UPLOAD_TRIGGERED.value, self.__vault_client.get_logger())
log_info(SkyflowMessages.Info.VALIDATING_FILE_UPLOAD_REQUEST.value, self.__vault_client.get_logger())
validate_file_upload_request(self.__vault_client.get_logger(), request)
diff --git a/skyflow/vault/data/__init__.py b/v2/skyflow/vault/data/__init__.py
similarity index 100%
rename from skyflow/vault/data/__init__.py
rename to v2/skyflow/vault/data/__init__.py
diff --git a/skyflow/vault/data/_delete_request.py b/v2/skyflow/vault/data/_delete_request.py
similarity index 100%
rename from skyflow/vault/data/_delete_request.py
rename to v2/skyflow/vault/data/_delete_request.py
diff --git a/skyflow/vault/data/_delete_response.py b/v2/skyflow/vault/data/_delete_response.py
similarity index 100%
rename from skyflow/vault/data/_delete_response.py
rename to v2/skyflow/vault/data/_delete_response.py
diff --git a/skyflow/vault/data/_file_upload_request.py b/v2/skyflow/vault/data/_file_upload_request.py
similarity index 100%
rename from skyflow/vault/data/_file_upload_request.py
rename to v2/skyflow/vault/data/_file_upload_request.py
diff --git a/skyflow/vault/data/_file_upload_response.py b/v2/skyflow/vault/data/_file_upload_response.py
similarity index 100%
rename from skyflow/vault/data/_file_upload_response.py
rename to v2/skyflow/vault/data/_file_upload_response.py
diff --git a/skyflow/vault/data/_get_request.py b/v2/skyflow/vault/data/_get_request.py
similarity index 100%
rename from skyflow/vault/data/_get_request.py
rename to v2/skyflow/vault/data/_get_request.py
diff --git a/skyflow/vault/data/_get_response.py b/v2/skyflow/vault/data/_get_response.py
similarity index 100%
rename from skyflow/vault/data/_get_response.py
rename to v2/skyflow/vault/data/_get_response.py
diff --git a/v2/skyflow/vault/data/_insert_request.py b/v2/skyflow/vault/data/_insert_request.py
new file mode 100644
index 00000000..2c6c2dcc
--- /dev/null
+++ b/v2/skyflow/vault/data/_insert_request.py
@@ -0,0 +1,20 @@
+from common.vault.data import BaseInsertRequest
+from skyflow.utils.enums import TokenMode
+
+class InsertRequest(BaseInsertRequest):
+ def __init__(self,
+ table: str,
+ values: list,
+ tokens: list = None,
+ upsert: str = None,
+ homogeneous: bool = False,
+ token_mode: TokenMode = TokenMode.DISABLE,
+ return_tokens: bool = True,
+ continue_on_error: bool = False):
+ super().__init__(table, values, upsert=upsert)
+ self.tokens = tokens
+ self.homogeneous = homogeneous
+ self.token_mode = token_mode
+ self.return_tokens = return_tokens
+ self.continue_on_error = continue_on_error
+
diff --git a/v2/skyflow/vault/data/_insert_response.py b/v2/skyflow/vault/data/_insert_response.py
new file mode 100644
index 00000000..0cf73343
--- /dev/null
+++ b/v2/skyflow/vault/data/_insert_response.py
@@ -0,0 +1,6 @@
+from common.vault.data import BaseInsertResponse
+
+
+class InsertResponse(BaseInsertResponse):
+ """PDB's own insert() response class -- currently identical to the shared base, kept as its
+ own subclass so PDB-specific fields can be added later without touching flowvault."""
diff --git a/skyflow/vault/data/_query_request.py b/v2/skyflow/vault/data/_query_request.py
similarity index 100%
rename from skyflow/vault/data/_query_request.py
rename to v2/skyflow/vault/data/_query_request.py
diff --git a/skyflow/vault/data/_query_response.py b/v2/skyflow/vault/data/_query_response.py
similarity index 100%
rename from skyflow/vault/data/_query_response.py
rename to v2/skyflow/vault/data/_query_response.py
diff --git a/skyflow/vault/data/_update_request.py b/v2/skyflow/vault/data/_update_request.py
similarity index 100%
rename from skyflow/vault/data/_update_request.py
rename to v2/skyflow/vault/data/_update_request.py
diff --git a/skyflow/vault/data/_update_response.py b/v2/skyflow/vault/data/_update_response.py
similarity index 100%
rename from skyflow/vault/data/_update_response.py
rename to v2/skyflow/vault/data/_update_response.py
diff --git a/skyflow/vault/data/_upload_file_request.py b/v2/skyflow/vault/data/_upload_file_request.py
similarity index 100%
rename from skyflow/vault/data/_upload_file_request.py
rename to v2/skyflow/vault/data/_upload_file_request.py
diff --git a/skyflow/vault/detect/__init__.py b/v2/skyflow/vault/detect/__init__.py
similarity index 100%
rename from skyflow/vault/detect/__init__.py
rename to v2/skyflow/vault/detect/__init__.py
diff --git a/skyflow/vault/detect/_audio_bleep.py b/v2/skyflow/vault/detect/_audio_bleep.py
similarity index 100%
rename from skyflow/vault/detect/_audio_bleep.py
rename to v2/skyflow/vault/detect/_audio_bleep.py
diff --git a/skyflow/vault/detect/_date_transformation.py b/v2/skyflow/vault/detect/_date_transformation.py
similarity index 100%
rename from skyflow/vault/detect/_date_transformation.py
rename to v2/skyflow/vault/detect/_date_transformation.py
diff --git a/skyflow/vault/detect/_deidentify_file_request.py b/v2/skyflow/vault/detect/_deidentify_file_request.py
similarity index 100%
rename from skyflow/vault/detect/_deidentify_file_request.py
rename to v2/skyflow/vault/detect/_deidentify_file_request.py
diff --git a/skyflow/vault/detect/_deidentify_file_response.py b/v2/skyflow/vault/detect/_deidentify_file_response.py
similarity index 100%
rename from skyflow/vault/detect/_deidentify_file_response.py
rename to v2/skyflow/vault/detect/_deidentify_file_response.py
diff --git a/skyflow/vault/detect/_deidentify_text_request.py b/v2/skyflow/vault/detect/_deidentify_text_request.py
similarity index 100%
rename from skyflow/vault/detect/_deidentify_text_request.py
rename to v2/skyflow/vault/detect/_deidentify_text_request.py
diff --git a/skyflow/vault/detect/_deidentify_text_response.py b/v2/skyflow/vault/detect/_deidentify_text_response.py
similarity index 100%
rename from skyflow/vault/detect/_deidentify_text_response.py
rename to v2/skyflow/vault/detect/_deidentify_text_response.py
diff --git a/skyflow/vault/detect/_entity_info.py b/v2/skyflow/vault/detect/_entity_info.py
similarity index 100%
rename from skyflow/vault/detect/_entity_info.py
rename to v2/skyflow/vault/detect/_entity_info.py
diff --git a/skyflow/vault/detect/_file.py b/v2/skyflow/vault/detect/_file.py
similarity index 100%
rename from skyflow/vault/detect/_file.py
rename to v2/skyflow/vault/detect/_file.py
diff --git a/skyflow/vault/detect/_file_input.py b/v2/skyflow/vault/detect/_file_input.py
similarity index 100%
rename from skyflow/vault/detect/_file_input.py
rename to v2/skyflow/vault/detect/_file_input.py
diff --git a/skyflow/vault/detect/_get_detect_run_request.py b/v2/skyflow/vault/detect/_get_detect_run_request.py
similarity index 100%
rename from skyflow/vault/detect/_get_detect_run_request.py
rename to v2/skyflow/vault/detect/_get_detect_run_request.py
diff --git a/skyflow/vault/detect/_reidentify_text_request.py b/v2/skyflow/vault/detect/_reidentify_text_request.py
similarity index 100%
rename from skyflow/vault/detect/_reidentify_text_request.py
rename to v2/skyflow/vault/detect/_reidentify_text_request.py
diff --git a/skyflow/vault/detect/_reidentify_text_response.py b/v2/skyflow/vault/detect/_reidentify_text_response.py
similarity index 100%
rename from skyflow/vault/detect/_reidentify_text_response.py
rename to v2/skyflow/vault/detect/_reidentify_text_response.py
diff --git a/skyflow/vault/detect/_text_index.py b/v2/skyflow/vault/detect/_text_index.py
similarity index 100%
rename from skyflow/vault/detect/_text_index.py
rename to v2/skyflow/vault/detect/_text_index.py
diff --git a/skyflow/vault/detect/_token_format.py b/v2/skyflow/vault/detect/_token_format.py
similarity index 100%
rename from skyflow/vault/detect/_token_format.py
rename to v2/skyflow/vault/detect/_token_format.py
diff --git a/skyflow/vault/detect/_transformations.py b/v2/skyflow/vault/detect/_transformations.py
similarity index 100%
rename from skyflow/vault/detect/_transformations.py
rename to v2/skyflow/vault/detect/_transformations.py
diff --git a/skyflow/vault/tokens/__init__.py b/v2/skyflow/vault/tokens/__init__.py
similarity index 100%
rename from skyflow/vault/tokens/__init__.py
rename to v2/skyflow/vault/tokens/__init__.py
diff --git a/skyflow/vault/tokens/_detokenize_request.py b/v2/skyflow/vault/tokens/_detokenize_request.py
similarity index 100%
rename from skyflow/vault/tokens/_detokenize_request.py
rename to v2/skyflow/vault/tokens/_detokenize_request.py
diff --git a/skyflow/vault/tokens/_detokenize_response.py b/v2/skyflow/vault/tokens/_detokenize_response.py
similarity index 100%
rename from skyflow/vault/tokens/_detokenize_response.py
rename to v2/skyflow/vault/tokens/_detokenize_response.py
diff --git a/skyflow/vault/tokens/_tokenize_request.py b/v2/skyflow/vault/tokens/_tokenize_request.py
similarity index 100%
rename from skyflow/vault/tokens/_tokenize_request.py
rename to v2/skyflow/vault/tokens/_tokenize_request.py
diff --git a/skyflow/vault/tokens/_tokenize_response.py b/v2/skyflow/vault/tokens/_tokenize_response.py
similarity index 100%
rename from skyflow/vault/tokens/_tokenize_response.py
rename to v2/skyflow/vault/tokens/_tokenize_response.py
diff --git a/v2/tests/__init__.py b/v2/tests/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/v2/tests/client/__init__.py b/v2/tests/client/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/client/test_skyflow.py b/v2/tests/client/test_skyflow.py
similarity index 91%
rename from tests/client/test_skyflow.py
rename to v2/tests/client/test_skyflow.py
index 5b7ea675..00d360fc 100644
--- a/tests/client/test_skyflow.py
+++ b/v2/tests/client/test_skyflow.py
@@ -124,7 +124,7 @@ def test_get_vault_with_invalid_vault_id_and_non_empty_list_raises_error(self):
SkyflowMessages.Error.VAULT_ID_NOT_IN_CONFIG_LIST.value.format("invalid_vault_id"),
)
- @patch("skyflow.client.skyflow.validate_vault_config")
+ @patch("skyflow.client.skyflow.Skyflow.Builder._validate_vault_config")
def test_build_calls_validate_vault_config(self, mock_validate_vault_config):
self.builder.add_vault_config(VALID_VAULT_CONFIG)
self.builder.build()
@@ -223,7 +223,7 @@ def test_get_connection_with_invalid_connection_id_and_empty_list_raises_Error(s
self.assertEqual(context.exception.message, SkyflowMessages.Error.EMPTY_CONNECTION_CONFIGS.value)
- @patch("skyflow.client.skyflow.validate_connection_config")
+ @patch("skyflow.client.skyflow.Skyflow.Builder._validate_connection_config")
def test_build_calls_validate_connection_config(self, mock_validate):
self.builder.add_connection_config(VALID_CONNECTION_CONFIG)
self.builder.build()
@@ -246,16 +246,16 @@ def test_invalid_credentials(self):
self.assertEqual(VALID_CREDENTIALS, self.builder._Builder__skyflow_credentials)
self.assertEqual(builder, self.builder)
- @patch("skyflow.client.skyflow.validate_vault_config")
+ @patch("skyflow.client.skyflow.Skyflow.Builder._validate_vault_config")
def test_skyflow_client_add_remove_vault_config(self, mock_validate_vault_config):
skyflow_client = self.builder.add_vault_config(VALID_VAULT_CONFIG).build()
new_config = VALID_VAULT_CONFIG.copy()
- new_config["vault_id"] = "VAULT_ID"
+ new_config["vault_id"] = "VAULT_ID_2"
skyflow_client.add_vault_config(new_config)
self.assertEqual(mock_validate_vault_config.call_count, 2)
- self.assertEqual("VAULT_ID", skyflow_client.get_vault_config(new_config["vault_id"]).get("vault_id"))
+ self.assertEqual("VAULT_ID_2", skyflow_client.get_vault_config(new_config["vault_id"]).get("vault_id"))
skyflow_client.remove_vault_config(new_config["vault_id"])
with self.assertRaises(SkyflowError) as context:
@@ -266,6 +266,16 @@ def test_skyflow_client_add_remove_vault_config(self, mock_validate_vault_config
SkyflowMessages.Error.VAULT_ID_NOT_IN_CONFIG_LIST.value.format(new_config["vault_id"]),
)
+ def test_skyflow_client_add_vault_config_duplicate_id_raises(self):
+ skyflow_client = self.builder.add_vault_config(VALID_VAULT_CONFIG).build()
+ with self.assertRaises(SkyflowError) as context:
+ skyflow_client.add_vault_config(VALID_VAULT_CONFIG)
+
+ self.assertEqual(
+ context.exception.message,
+ SkyflowMessages.Error.VAULT_ID_ALREADY_EXISTS.value.format(VALID_VAULT_CONFIG["vault_id"]),
+ )
+
@patch("skyflow.vault.client.client.VaultClient.update_config")
def test_skyflow_client_update_and_get_vault_config(self, mock_update_config):
skyflow_client = self.builder.add_vault_config(VALID_VAULT_CONFIG).build()
@@ -278,19 +288,19 @@ def test_skyflow_client_update_and_get_vault_config(self, mock_update_config):
self.assertEqual(VALID_VAULT_CONFIG.get("vault_id"), vault.get("vault_id"))
- @patch("skyflow.client.skyflow.validate_connection_config")
+ @patch("skyflow.client.skyflow.Skyflow.Builder._validate_connection_config")
def test_skyflow_client_add_remove_connection_config(self, mock_validate_connection_config):
skyflow_client = self.builder.add_connection_config(VALID_CONNECTION_CONFIG).build()
new_config = VALID_CONNECTION_CONFIG.copy()
- new_config["connection_id"] = "CONNECTION_ID"
+ new_config["connection_id"] = "CONNECTION_ID_2"
skyflow_client.add_connection_config(new_config)
self.assertEqual(mock_validate_connection_config.call_count, 2)
self.assertEqual(
- "CONNECTION_ID", skyflow_client.get_connection_config(new_config["connection_id"]).get("connection_id")
+ "CONNECTION_ID_2", skyflow_client.get_connection_config(new_config["connection_id"]).get("connection_id")
)
- skyflow_client.remove_connection_config("CONNECTION_ID")
+ skyflow_client.remove_connection_config("CONNECTION_ID_2")
with self.assertRaises(SkyflowError) as context:
skyflow_client.get_connection_config(new_config["connection_id"]).get("connection_id")
@@ -299,6 +309,16 @@ def test_skyflow_client_add_remove_connection_config(self, mock_validate_connect
SkyflowMessages.Error.CONNECTION_ID_NOT_IN_CONFIG_LIST.value.format(new_config["connection_id"]),
)
+ def test_skyflow_client_add_connection_config_duplicate_id_raises(self):
+ skyflow_client = self.builder.add_connection_config(VALID_CONNECTION_CONFIG).build()
+ with self.assertRaises(SkyflowError) as context:
+ skyflow_client.add_connection_config(VALID_CONNECTION_CONFIG)
+
+ self.assertEqual(
+ context.exception.message,
+ SkyflowMessages.Error.CONNECTION_ID_ALREADY_EXISTS.value.format(VALID_CONNECTION_CONFIG["connection_id"]),
+ )
+
@patch("skyflow.vault.client.client.VaultClient.update_config")
def test_skyflow_client_update_and_get_connection_config(self, mock_update_config):
builder = self.builder
@@ -392,21 +412,21 @@ def test_update_connection_config_with_invalid_connection_id_raises_error(self,
class TestVaultClient(unittest.TestCase):
def _make_client(self):
client = VaultClient({"vault_id": "test_vault"})
- client._VaultClient__api_client = Mock()
+ client._api_client = Mock()
return client
def test_get_detect_text_api_returns_strings(self):
client = self._make_client()
result = client.get_detect_text_api()
- self.assertEqual(result, client._VaultClient__api_client.strings)
+ self.assertEqual(result, client._api_client.strings)
def test_get_detect_file_api_returns_files(self):
client = self._make_client()
result = client.get_detect_file_api()
- self.assertEqual(result, client._VaultClient__api_client.files)
+ self.assertEqual(result, client._api_client.files)
- @patch("skyflow.vault.client.client.generate_bearer_token_from_creds")
- @patch("skyflow.vault.client.client.is_expired", return_value=True)
+ @patch("common.vault.base_vault_client.generate_bearer_token_from_creds")
+ @patch("common.vault.base_vault_client.is_expired", return_value=True)
def test_get_bearer_token_passes_token_uri_option(self, _mock_expired, mock_gen):
mock_gen.return_value = ("test_token", "bearer")
client = VaultClient({"vault_id": "test_vault"})
@@ -426,7 +446,7 @@ def _build_client(self):
def test_update_log_level_emits_deprecation_warning(self):
client = self._build_client()
- with patch('skyflow.client.skyflow.log_warn') as mock_warn:
+ with patch('common.client.base_skyflow.log_warn') as mock_warn:
client.update_log_level(LogLevel.INFO)
mock_warn.assert_called_once()
self.assertIn("set_log_level", mock_warn.call_args[0][0])
diff --git a/v2/tests/service_account/__init__.py b/v2/tests/service_account/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/service_account/invalid_creds.json b/v2/tests/service_account/invalid_creds.json
similarity index 100%
rename from tests/service_account/invalid_creds.json
rename to v2/tests/service_account/invalid_creds.json
diff --git a/tests/service_account/test__utils.py b/v2/tests/service_account/test__utils.py
similarity index 100%
rename from tests/service_account/test__utils.py
rename to v2/tests/service_account/test__utils.py
diff --git a/v2/tests/utils/__init__.py b/v2/tests/utils/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/v2/tests/utils/logger/__init__.py b/v2/tests/utils/logger/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/utils/logger/test__log_helpers.py b/v2/tests/utils/logger/test__log_helpers.py
similarity index 87%
rename from tests/utils/logger/test__log_helpers.py
rename to v2/tests/utils/logger/test__log_helpers.py
index 1ea50d45..7cf2e9db 100644
--- a/tests/utils/logger/test__log_helpers.py
+++ b/v2/tests/utils/logger/test__log_helpers.py
@@ -7,7 +7,7 @@
class TestLoggingFunctions(unittest.TestCase):
- @patch('skyflow.utils.logger._log_helpers.Logger')
+ @patch('common.utils.logger._log_helpers.Logger')
def test_log_info_with_logger(self, MockLogger):
mock_logger = MockLogger()
message = "Info message"
@@ -17,14 +17,14 @@ def test_log_info_with_logger(self, MockLogger):
mock_logger.info.assert_called_once_with(f"{message}")
- @patch('skyflow.utils.logger._log_helpers.Logger')
+ @patch('common.utils.logger._log_helpers.Logger')
def test_log_info_without_logger(self, MockLogger):
try:
log_info("Message", None)
except AttributeError:
self.fail("log_info raised AttributeError unexpectedly!")
- @patch('skyflow.utils.logger._log_helpers.Logger')
+ @patch('common.utils.logger._log_helpers.Logger')
def test_log_error_with_all_fields(self, MockLogger):
mock_logger = MockLogger()
message = "Error message"
@@ -47,7 +47,7 @@ def test_log_error_with_all_fields(self, MockLogger):
mock_logger.error.assert_called_once_with(expected_log_data)
- @patch('skyflow.utils.logger._log_helpers.Logger')
+ @patch('common.utils.logger._log_helpers.Logger')
def test_log_error_with_minimal_fields(self, MockLogger):
mock_logger = MockLogger()
message = "Minimal error"
@@ -62,7 +62,7 @@ def test_log_error_with_minimal_fields(self, MockLogger):
mock_logger.error.assert_called_once_with(expected_log_data)
- @patch('skyflow.utils.logger._log_helpers.Logger')
+ @patch('common.utils.logger._log_helpers.Logger')
def test_log_error_creates_logger_if_none(self, MockLogger):
message = "Auto-created logger error"
http_code = 500
@@ -71,7 +71,7 @@ def test_log_error_creates_logger_if_none(self, MockLogger):
MockLogger.assert_called_once_with(LogLevel.ERROR)
- @patch('skyflow.utils.logger._log_helpers.Logger')
+ @patch('common.utils.logger._log_helpers.Logger')
def test_log_error_handles_missing_optional_fields(self, MockLogger):
mock_logger = MockLogger()
message = "Test missing optional fields"
diff --git a/tests/utils/logger/test__logger.py b/v2/tests/utils/logger/test__logger.py
similarity index 100%
rename from tests/utils/logger/test__logger.py
rename to v2/tests/utils/logger/test__logger.py
diff --git a/tests/utils/test__helpers.py b/v2/tests/utils/test__helpers.py
similarity index 100%
rename from tests/utils/test__helpers.py
rename to v2/tests/utils/test__helpers.py
diff --git a/tests/utils/test__utils.py b/v2/tests/utils/test__utils.py
similarity index 100%
rename from tests/utils/test__utils.py
rename to v2/tests/utils/test__utils.py
diff --git a/v2/tests/utils/validations/__init__.py b/v2/tests/utils/validations/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/utils/validations/test__validations.py b/v2/tests/utils/validations/test__validations.py
similarity index 100%
rename from tests/utils/validations/test__validations.py
rename to v2/tests/utils/validations/test__validations.py
diff --git a/v2/tests/vault/__init__.py b/v2/tests/vault/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/v2/tests/vault/client/__init__.py b/v2/tests/vault/client/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/v2/tests/vault/client/test__client.py b/v2/tests/vault/client/test__client.py
new file mode 100644
index 00000000..a0894567
--- /dev/null
+++ b/v2/tests/vault/client/test__client.py
@@ -0,0 +1,127 @@
+import unittest
+from unittest.mock import patch, MagicMock
+
+from common.utils.enums import Env
+from skyflow.error import SkyflowError
+from skyflow.utils import SkyflowMessages
+from skyflow.vault.client.client import VaultClient
+
+CONFIG = {
+ "credentials": "some_credentials",
+ "cluster_id": "test_cluster_id",
+ "env": "test_env",
+ "vault_id": "test_vault_id",
+ "roles": ["role_id_1", "role_id_2"],
+ "ctx": "context"
+}
+
+CREDENTIALS_WITH_API_KEY = {"api_key": "dummy_api_key"}
+CREDENTIALS_WITH_TOKEN = {"token": "dummy_static_token"}
+CREDENTIALS_WITH_PATH = {"path": "/some/path/credentials.json"}
+CREDENTIALS_WITH_STRING = {"credentials_string": '{"clientID": "x"}'}
+
+
+class TestVaultClient(unittest.TestCase):
+ """v2-specific VaultClient coverage.
+
+ Shared logic (initialize_client_configuration, get_bearer_token, credential-resolution
+ fast paths, update_config) moved to common.vault.base_vault_client.BaseVaultClient and is
+ covered by common/tests/vault/test_base_vault_client.py instead -- those scenarios used to
+ live here, patching names at this module's path, but the code they exercise no longer runs
+ from this module. What remains here is what's genuinely still v2-specific: the
+ initialize_api_client() lambda/Skyflow-construction override, and the resource accessors.
+ """
+
+ def setUp(self):
+ self.vault_client = VaultClient(CONFIG)
+
+ # ------------------------------------------------------------------ #
+ # Basic setters / getters (inherited from BaseVaultClient, but exercised
+ # here too as a cheap smoke test that the subclass wiring works)
+ # ------------------------------------------------------------------ #
+
+ def test_set_common_skyflow_credentials(self):
+ credentials = {"api_key": "dummy_api_key"}
+ self.vault_client.set_common_skyflow_credentials(credentials)
+ self.assertEqual(self.vault_client.get_common_skyflow_credentials(), credentials)
+
+ def test_set_logger(self):
+ mock_logger = MagicMock()
+ self.vault_client.set_logger("INFO", mock_logger)
+ self.assertEqual(self.vault_client.get_log_level(), "INFO")
+ self.assertEqual(self.vault_client.get_logger(), mock_logger)
+
+ def test_get_vault_id(self):
+ self.assertEqual(self.vault_client.get_vault_id(), CONFIG["vault_id"])
+
+ def test_get_config(self):
+ self.assertEqual(self.vault_client.get_config(), CONFIG)
+
+ # ------------------------------------------------------------------ #
+ # resolve_vault_url — v2's own domain (vault.skyflowapis.*), the OTHER
+ # hook v2 overrides. Regression-pins the exact host v2 must keep hitting.
+ # ------------------------------------------------------------------ #
+
+ def test_resolve_vault_url_uses_v2_domain(self):
+ url = self.vault_client.resolve_vault_url("mycluster", Env.PROD, "myvault")
+ self.assertEqual(url, "https://mycluster.vault.skyflowapis.com")
+
+ # ------------------------------------------------------------------ #
+ # initialize_api_client — lambda token provider (v2-specific: this is
+ # the one hook v2 actually overrides)
+ # ------------------------------------------------------------------ #
+
+ @patch("skyflow.vault.client.client.Skyflow")
+ def test_initialize_api_client_passes_callable_token(self, mock_skyflow):
+ """initialize_api_client must pass a callable (lambda) as token, not a string."""
+ self.vault_client.initialize_api_client("https://test-vault-url.com", "initial_token")
+
+ args, kwargs = mock_skyflow.call_args
+ self.assertEqual(kwargs["base_url"], "https://test-vault-url.com")
+ self.assertTrue(callable(kwargs["token"]), "token must be a callable (lambda)")
+
+ @patch("skyflow.vault.client.client.Skyflow")
+ def test_initialize_api_client_lambda_returns_cached_bearer_token(self, mock_skyflow):
+ """Lambda returns _bearer_token when it is set (interceptor behaviour)."""
+ self.vault_client._bearer_token = "refreshed_token"
+ self.vault_client.initialize_api_client("https://test-vault-url.com", "initial_token")
+
+ _, kwargs = mock_skyflow.call_args
+ self.assertEqual(kwargs["token"](), "refreshed_token")
+
+ @patch("skyflow.vault.client.client.Skyflow")
+ def test_initialize_api_client_lambda_falls_back_to_initial_token(self, mock_skyflow):
+ """Lambda falls back to the initial token when _bearer_token is None."""
+ self.vault_client._bearer_token = None
+ self.vault_client.initialize_api_client("https://test-vault-url.com", "initial_token")
+
+ _, kwargs = mock_skyflow.call_args
+ self.assertEqual(kwargs["token"](), "initial_token")
+
+ # ------------------------------------------------------------------ #
+ # API accessor stubs (v2-specific: v3 doesn't have these)
+ # ------------------------------------------------------------------ #
+
+ def test_get_records_api(self):
+ self.vault_client._api_client = MagicMock()
+ self.assertIsNotNone(self.vault_client.get_records_api())
+
+ def test_get_tokens_api(self):
+ self.vault_client._api_client = MagicMock()
+ self.assertIsNotNone(self.vault_client.get_tokens_api())
+
+ def test_get_query_api(self):
+ self.vault_client._api_client = MagicMock()
+ self.assertIsNotNone(self.vault_client.get_query_api())
+
+ def test_get_detect_text_api(self):
+ self.vault_client._api_client = MagicMock()
+ self.assertIsNotNone(self.vault_client.get_detect_text_api())
+
+ def test_get_detect_file_api(self):
+ self.vault_client._api_client = MagicMock()
+ self.assertIsNotNone(self.vault_client.get_detect_file_api())
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/v2/tests/vault/connection/__init__.py b/v2/tests/vault/connection/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/vault/connection/test_responses.py b/v2/tests/vault/connection/test_responses.py
similarity index 100%
rename from tests/vault/connection/test_responses.py
rename to v2/tests/vault/connection/test_responses.py
diff --git a/v2/tests/vault/controller/__init__.py b/v2/tests/vault/controller/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/vault/controller/test__audit_binlookup.py b/v2/tests/vault/controller/test__audit_binlookup.py
similarity index 100%
rename from tests/vault/controller/test__audit_binlookup.py
rename to v2/tests/vault/controller/test__audit_binlookup.py
diff --git a/tests/vault/controller/test__connection.py b/v2/tests/vault/controller/test__connection.py
similarity index 100%
rename from tests/vault/controller/test__connection.py
rename to v2/tests/vault/controller/test__connection.py
diff --git a/tests/vault/controller/test__detect.py b/v2/tests/vault/controller/test__detect.py
similarity index 100%
rename from tests/vault/controller/test__detect.py
rename to v2/tests/vault/controller/test__detect.py
diff --git a/tests/vault/controller/test__vault.py b/v2/tests/vault/controller/test__vault.py
similarity index 96%
rename from tests/vault/controller/test__vault.py
rename to v2/tests/vault/controller/test__vault.py
index 5acdf779..f3c9124c 100644
--- a/tests/vault/controller/test__vault.py
+++ b/v2/tests/vault/controller/test__vault.py
@@ -140,6 +140,31 @@ def test_insert_handles_generic_error(self, mock_validate):
records_api.with_raw_response.record_service_insert_record.assert_called_once()
+ @patch("skyflow.vault.controller._vault.validate_insert_request")
+ def test_insert_raises_on_empty_key(self, mock_validate):
+ """Shared BaseVaultController._validate_field_values() -- an empty/null key in insert
+ data now raises for v2 too (previously only logged a warning and let the insert through)."""
+ request = InsertRequest(table="test_table", values=[{"": "value"}])
+
+ with self.assertRaises(SkyflowError):
+ self.vault.insert(request)
+
+ self.vault_client.get_records_api.return_value.with_raw_response.record_service_insert_record.assert_not_called()
+
+ @patch("skyflow.vault.controller._vault.validate_insert_request")
+ @patch("skyflow.vault.controller._vault.parse_insert_response")
+ def test_insert_allows_empty_value(self, mock_parse_response, mock_validate):
+ request = InsertRequest(table="test_table", values=[{"column_name": ""}])
+
+ mock_api_response = Mock()
+ mock_parse_response.return_value = InsertResponse(inserted_fields=[{"skyflow_id": "id1"}])
+ records_api = self.vault_client.get_records_api.return_value
+ records_api.with_raw_response.record_service_insert_record.return_value = mock_api_response
+
+ self.vault.insert(request)
+
+ records_api.with_raw_response.record_service_insert_record.assert_called_once()
+
@patch("skyflow.vault.controller._vault.validate_insert_request")
@patch("skyflow.vault.controller._vault.parse_insert_response")
def test_insert_with_continue_on_error_false_when_tokens_are_not_none(self, mock_parse_response, mock_validate):
diff --git a/v2/tests/vault/data/__init__.py b/v2/tests/vault/data/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/vault/data/test_responses.py b/v2/tests/vault/data/test_responses.py
similarity index 93%
rename from tests/vault/data/test_responses.py
rename to v2/tests/vault/data/test_responses.py
index ea9f2be1..67616fbd 100644
--- a/tests/vault/data/test_responses.py
+++ b/v2/tests/vault/data/test_responses.py
@@ -1,4 +1,5 @@
import unittest
+from common.vault.data import BaseInsertResponse
from skyflow.vault.data._delete_response import DeleteResponse
from skyflow.vault.data._file_upload_response import FileUploadResponse
from skyflow.vault.data._get_response import GetResponse
@@ -54,6 +55,10 @@ def test_empty_data_not_replaced(self):
class TestInsertResponse(unittest.TestCase):
+ def test_is_a_base_insert_response(self):
+ r = InsertResponse(inserted_fields=[{"skyflow_id": "id1"}], errors=None)
+ self.assertIsInstance(r, BaseInsertResponse)
+
def test_repr(self):
r = InsertResponse(inserted_fields=[{"skyflow_id": "id1"}], errors=None)
self.assertIn("InsertResponse", repr(r))
diff --git a/v2/tests/vault/detect/__init__.py b/v2/tests/vault/detect/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/vault/detect/test_models.py b/v2/tests/vault/detect/test_models.py
similarity index 100%
rename from tests/vault/detect/test_models.py
rename to v2/tests/vault/detect/test_models.py
diff --git a/v2/tests/vault/tokens/__init__.py b/v2/tests/vault/tokens/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/vault/tokens/test_responses.py b/v2/tests/vault/tokens/test_responses.py
similarity index 100%
rename from tests/vault/tokens/test_responses.py
rename to v2/tests/vault/tokens/test_responses.py