diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f5c7d64..9545916 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -7,7 +7,7 @@ on: pull_request: env: - DEFAULT_PYTHON: "3.12" + DEFAULT_PYTHON: "3.13" DEFAULT_OS: ubuntu-latest jobs: @@ -30,14 +30,9 @@ jobs: run: | python -m pip install --upgrade pip python -m pip install ".[dev,spark]" - wget https://dlcdn.apache.org/spark/spark-4.0.1/spark-4.0.1-bin-hadoop3.tgz - tar -xzf spark-4.0.1-bin-hadoop3.tgz - export SPARK_HOME=$(pwd)/spark-4.0.1-bin-hadoop3 - export PATH=$SPARK_HOME/sbin:$PATH - start-thriftserver.sh - name: Run pytest with coverage run: | - CHRONIFY_HIVE_URL=hive://localhost:10000/default pytest -v --cov --cov-report=xml + pytest -v --cov --cov-report=xml - name: codecov uses: codecov/codecov-action@v4.2.0 if: ${{ matrix.os == env.DEFAULT_OS && matrix.python-version == env.DEFAULT_PYTHON }} @@ -53,11 +48,11 @@ jobs: - name: Set up Python uses: actions/setup-python@v5 with: - python-version: 3.12 + python-version: 3.13 - name: Install dependencies run: | python -m pip install --upgrade pip - python -m pip install ".[dev]" + python -m pip install ".[dev,pyspark]" mypy ruff: runs-on: ubuntu-latest diff --git a/docs/how_tos/spark_backend.md b/docs/how_tos/spark_backend.md index 1df45c6..c586ead 100644 --- a/docs/how_tos/spark_backend.md +++ b/docs/how_tos/spark_backend.md @@ -1,48 +1,52 @@ # Apache Spark Backend + +This guide shows how to use Chronify with Apache Spark for processing large time series datasets +that are too large for DuckDB on a single machine. + +## Prerequisites + Download Spark from https://spark.apache.org/downloads.html and install it. Spark provides startup scripts for UNIX operating systems (not Windows). ## Install chronify with Spark support -``` -$ pip install chronify --group=pyhive +```bash +pip install chronify[spark] ``` ## Installation on a development computer -Installation can be as simple as -``` -$ tar -xzf spark-4.0.1-bin-hadoop3.tgz -$ export SPARK_HOME=$(pwd)/spark-4.0.1-bin-hadoop3 -``` -Start a Thrift server. This allows JDBC clients to send SQL queries to an in-process Spark cluster -running in local mode. -``` -$ $SPARK_HOME/sbin/start-thriftserver.sh --master=spark://$(hostname):7077 +Installation can be as simple as: +```bash +tar -xzf spark-4.0.0-bin-hadoop3.tgz +export SPARK_HOME=$(pwd)/spark-4.0.0-bin-hadoop3 ``` -The URL to connect to this server is `hive://localhost:10000/default` - ## Installation on an HPC -The chronify development team uses these -[scripts](https://github.com/NREL/HPC/tree/master/applications/spark) to run Spark on NREL's HPC. +The chronify development team uses the +[package](https://github.com/NREL/sparkctl) to run Spark on NLR's HPC. ## Chronify Usage + This example creates a chronify Store with Spark as the backend and then adds a view to a Parquet file. Chronify will run its normal time checks. -First, create the Parquet file and chronify schema. +**Note:** The Spark backend is designed primarily for reading data from Parquet files via views. +Direct data ingestion (inserting DataFrames) is not supported with Spark. + +### Step 1: Create a Parquet file and schema ```python from datetime import datetime, timedelta import numpy as np import pandas as pd -from chronify import DatetimeRange, Store, TableSchema, CsvTableSchema +from chronify import DatetimeRange, Store, TableSchema initial_time = datetime(2020, 1, 1) end_time = datetime(2020, 12, 31, 23) resolution = timedelta(hours=1) timestamps = pd.date_range(initial_time, end_time, freq=resolution, unit="us") + dfs = [] for i in range(1, 4): df = pd.DataFrame( @@ -55,6 +59,7 @@ for i in range(1, 4): dfs.append(df) df = pd.concat(dfs) df.to_parquet("data.parquet", index=False) + schema = TableSchema( name="devices", value_column="value", @@ -68,14 +73,25 @@ schema = TableSchema( ) ``` +### Step 2: Create a Spark store and add a view + ```python from chronify import Store -store = Store.create_new_hive_store("hive://localhost:10000/default") -store.create_view_from_parquet("data.parquet") +# Create a new Spark store (this will create a local SparkSession) +store = Store.create_new_spark_store() + +# Or pass an existing SparkSession: +# from pyspark.sql import SparkSession +# spark = SparkSession.builder.appName("chronify").getOrCreate() +# store = Store.create_new_spark_store(session=spark) + +# Create a view from the Parquet file +store.create_view_from_parquet("data.parquet", schema) ``` -Verify the data: +### Step 3: Verify the data + ```python store.read_table(schema.name).head() ``` @@ -88,9 +104,94 @@ store.read_table(schema.name).head() 4 2020-01-01 04:00:00 1 0.994851 ``` +## Registering Ibis tables + +If you're using Ibis to load and transform data before passing it to chronify for time validation, +you must use `create_view`. This is useful when integrating with +other packages like dsgrid that use Ibis for data loading and transformation. + +### Validation only (recommended for most workflows) + +Use `create_view` when you only need time validation and don't need the data to persist +in chronify after the session ends: + +```python +import ibis +from chronify import Store, TableSchema, DatetimeRange +from datetime import datetime, timedelta + +# Load data with Ibis (can be from any Ibis-supported backend) +conn = ibis.pyspark.connect(spark_session) +table = conn.read_parquet("data.parquet") + +# Optionally transform the data (rename columns, filter, etc.) +table = table.rename({"old_timestamp": "timestamp"}) + +# Create the chronify schema +schema = TableSchema( + name="devices", + value_column="value", + time_config=DatetimeRange( + time_column="timestamp", + start=datetime(2020, 1, 1), + length=8784, + resolution=timedelta(hours=1), + ), + time_array_id_columns=["id"], +) + +# Register with chronify for time validation only +store = Store.create_new_spark_store(session=spark_session) +store.create_view(schema, table) +``` + +Both the temp view and schema are registered and available during the session. + +### Persistent registration (requires Hive) + +Use `create_view` when you need the view and schema to persist across sessions. +This requires Spark to be configured with Hive metastore support. + +#### Configuring Hive metastore + +To enable Hive support, create your SparkSession with `enableHiveSupport()`: + +```python +from pyspark.sql import SparkSession + +spark = SparkSession.builder \ + .appName("chronify") \ + .config("spark.sql.warehouse.dir", "/path/to/spark-warehouse") \ + .enableHiveSupport() \ + .getOrCreate() + +store = Store.create_new_spark_store(session=spark) +``` + ## Time configuration mapping + The primary use case for Spark is to map datasets that are larger than can be processed by DuckDB -on one computer. In such a workflow a user would call +on one computer. In such a workflow a user would call: + ```python store.map_table_time_config(src_table_name, dst_schema, output_file="mapped_data.parquet") ``` + +This writes the mapped data to a Parquet file that can then be used with any backend. + +## Backend differences + +The Spark backend differs from DuckDB and SQLite in the following ways: + +| Feature | DuckDB/SQLite | Spark | +|---------|---------------|-------| +| Data ingestion (`ingest_table`) | Yes | No | +| View from Parquet | Yes | Yes | +| Register Ibis table (`create_view`) | Yes | Yes | +| Delete rows | Yes | No | +| In-memory storage | Yes | No (views only) | +| Time mapping to Parquet | Yes | Yes | + +For smaller datasets that fit in memory, DuckDB offers better performance and full feature support. +Use Spark when working with datasets too large for a single machine or when integrating with +existing Spark infrastructure. diff --git a/pyproject.toml b/pyproject.toml index bab7845..5f70d3d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,24 +27,20 @@ classifiers = [ ] dependencies = [ "duckdb ~= 1.1.0", - "duckdb_engine", + "ibis-framework[duckdb,sqlite] >= 9.0", "loguru", "pandas >= 2.2, < 3", "pyarrow", "pydantic >= 2.7, < 3", "pytz", "rich", - "sqlalchemy == 2.0.37", "tzdata", - # Required by pyhive - "future", - "python-dateutil", ] [project.optional-dependencies] spark = [ - "thrift", - "thrift_sasl", + "ibis-framework[pyspark]", + "pyspark == 4.0.0", ] dev = [ @@ -63,11 +59,6 @@ dev = [ "sphinx-tabs~=3.4", ] -[project.entry-points."sqlalchemy.dialects"] -hive = "pyhive.sqlalchemy_hive:HiveDialect" -"hive.http" = "pyhive.sqlalchemy_hive:HiveHTTPDialect" -"hive.https" = "pyhive.sqlalchemy_hive:HiveHTTPSDialect" - [project.urls] Documentation = "https://github.com/NREL/chronify#readme" Issues = "https://github.com/NREL/chronify/issues" @@ -77,11 +68,16 @@ Source = "https://github.com/NREL/chronify" files = [ "src", ] -exclude = [ - "src/chronify/_vendor/*", -] strict = true +[[tool.mypy.overrides]] +module = [ + "ibis.*", + "pyarrow.*", + "pyspark.*", +] +ignore_missing_imports = true + [tool.pytest.ini_options] pythonpath = "src" minversion = "6.0" @@ -99,7 +95,6 @@ exclude = [ "dist", "env", "venv", - "src/chronify/_vendor/*", ] line-length = 99 diff --git a/src/chronify/__init__.py b/src/chronify/__init__.py index 4637d5d..02128f2 100644 --- a/src/chronify/__init__.py +++ b/src/chronify/__init__.py @@ -1,8 +1,5 @@ import importlib.metadata as metadata -import sys -from chronify._vendor.kyuubi import TCLIService -from chronify._vendor.kyuubi import pyhive from chronify.exceptions import ( ChronifyExceptionBase, ConflictingInputsError, @@ -60,8 +57,3 @@ ) __version__ = metadata.metadata("chronify")["Version"] - - -# Make pyhive importable as if it were installed separately. -sys.modules["pyhive"] = pyhive -sys.modules["TCLIService"] = TCLIService diff --git a/src/chronify/_vendor/kyuubi/LICENSE b/src/chronify/_vendor/kyuubi/LICENSE deleted file mode 100644 index f49a4e1..0000000 --- a/src/chronify/_vendor/kyuubi/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. \ No newline at end of file diff --git a/src/chronify/_vendor/kyuubi/README.md b/src/chronify/_vendor/kyuubi/README.md deleted file mode 100644 index b71db55..0000000 --- a/src/chronify/_vendor/kyuubi/README.md +++ /dev/null @@ -1,12 +0,0 @@ -This source code is copied from https://github.com/apache/kyuubi.git, -commit ID 3b205a3924e0e3a75c425de1396089729cf22ee5. We did not modify the code. - -The pyhive package is marked as not supported, but we need it to work with Spark. -The latest published version of pyhive on pypi.org is not compatible with current versions of -sqlalchemy and Apache Spark. Specifically, we require the patch made in commit ID -a0b9873f817267675eab304f6935bafa4ab0f731. - -We have validated this version of pyhive with our use cases. We will remove this code as -soon as Kyuubi publishes an updated version on pypi.org. - -The pyhive license file is included here. diff --git a/src/chronify/_vendor/kyuubi/TCLIService/TCLIService-remote b/src/chronify/_vendor/kyuubi/TCLIService/TCLIService-remote deleted file mode 100755 index 8d875fa..0000000 --- a/src/chronify/_vendor/kyuubi/TCLIService/TCLIService-remote +++ /dev/null @@ -1,264 +0,0 @@ -#!/usr/bin/env python -# -# Autogenerated by Thrift Compiler (0.10.0) -# -# DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING -# -# options string: py -# - -import sys -import pprint -if sys.version_info[0] > 2: - from urllib.parse import urlparse -else: - from urlparse import urlparse -from thrift.transport import TTransport, TSocket, TSSLSocket, THttpClient -from thrift.protocol.TBinaryProtocol import TBinaryProtocol - -from TCLIService import TCLIService -from TCLIService.ttypes import * - -if len(sys.argv) <= 1 or sys.argv[1] == '--help': - print('') - print('Usage: ' + sys.argv[0] + ' [-h host[:port]] [-u url] [-f[ramed]] [-s[sl]] [-novalidate] [-ca_certs certs] [-keyfile keyfile] [-certfile certfile] function [arg1 [arg2...]]') - print('') - print('Functions:') - print(' TOpenSessionResp OpenSession(TOpenSessionReq req)') - print(' TCloseSessionResp CloseSession(TCloseSessionReq req)') - print(' TGetInfoResp GetInfo(TGetInfoReq req)') - print(' TExecuteStatementResp ExecuteStatement(TExecuteStatementReq req)') - print(' TGetTypeInfoResp GetTypeInfo(TGetTypeInfoReq req)') - print(' TGetCatalogsResp GetCatalogs(TGetCatalogsReq req)') - print(' TGetSchemasResp GetSchemas(TGetSchemasReq req)') - print(' TGetTablesResp GetTables(TGetTablesReq req)') - print(' TGetTableTypesResp GetTableTypes(TGetTableTypesReq req)') - print(' TGetColumnsResp GetColumns(TGetColumnsReq req)') - print(' TGetFunctionsResp GetFunctions(TGetFunctionsReq req)') - print(' TGetPrimaryKeysResp GetPrimaryKeys(TGetPrimaryKeysReq req)') - print(' TGetCrossReferenceResp GetCrossReference(TGetCrossReferenceReq req)') - print(' TGetOperationStatusResp GetOperationStatus(TGetOperationStatusReq req)') - print(' TCancelOperationResp CancelOperation(TCancelOperationReq req)') - print(' TCloseOperationResp CloseOperation(TCloseOperationReq req)') - print(' TGetResultSetMetadataResp GetResultSetMetadata(TGetResultSetMetadataReq req)') - print(' TFetchResultsResp FetchResults(TFetchResultsReq req)') - print(' TGetDelegationTokenResp GetDelegationToken(TGetDelegationTokenReq req)') - print(' TCancelDelegationTokenResp CancelDelegationToken(TCancelDelegationTokenReq req)') - print(' TRenewDelegationTokenResp RenewDelegationToken(TRenewDelegationTokenReq req)') - print(' TGetLogResp GetLog(TGetLogReq req)') - print('') - sys.exit(0) - -pp = pprint.PrettyPrinter(indent=2) -host = 'localhost' -port = 9090 -uri = '' -framed = False -ssl = False -validate = True -ca_certs = None -keyfile = None -certfile = None -http = False -argi = 1 - -if sys.argv[argi] == '-h': - parts = sys.argv[argi + 1].split(':') - host = parts[0] - if len(parts) > 1: - port = int(parts[1]) - argi += 2 - -if sys.argv[argi] == '-u': - url = urlparse(sys.argv[argi + 1]) - parts = url[1].split(':') - host = parts[0] - if len(parts) > 1: - port = int(parts[1]) - else: - port = 80 - uri = url[2] - if url[4]: - uri += '?%s' % url[4] - http = True - argi += 2 - -if sys.argv[argi] == '-f' or sys.argv[argi] == '-framed': - framed = True - argi += 1 - -if sys.argv[argi] == '-s' or sys.argv[argi] == '-ssl': - ssl = True - argi += 1 - -if sys.argv[argi] == '-novalidate': - validate = False - argi += 1 - -if sys.argv[argi] == '-ca_certs': - ca_certs = sys.argv[argi+1] - argi += 2 - -if sys.argv[argi] == '-keyfile': - keyfile = sys.argv[argi+1] - argi += 2 - -if sys.argv[argi] == '-certfile': - certfile = sys.argv[argi+1] - argi += 2 - -cmd = sys.argv[argi] -args = sys.argv[argi + 1:] - -if http: - transport = THttpClient.THttpClient(host, port, uri) -else: - if ssl: - socket = TSSLSocket.TSSLSocket(host, port, validate=validate, ca_certs=ca_certs, keyfile=keyfile, certfile=certfile) - else: - socket = TSocket.TSocket(host, port) - if framed: - transport = TTransport.TFramedTransport(socket) - else: - transport = TTransport.TBufferedTransport(socket) -protocol = TBinaryProtocol(transport) -client = TCLIService.Client(protocol) -transport.open() - -if cmd == 'OpenSession': - if len(args) != 1: - print('OpenSession requires 1 args') - sys.exit(1) - pp.pprint(client.OpenSession(eval(args[0]),)) - -elif cmd == 'CloseSession': - if len(args) != 1: - print('CloseSession requires 1 args') - sys.exit(1) - pp.pprint(client.CloseSession(eval(args[0]),)) - -elif cmd == 'GetInfo': - if len(args) != 1: - print('GetInfo requires 1 args') - sys.exit(1) - pp.pprint(client.GetInfo(eval(args[0]),)) - -elif cmd == 'ExecuteStatement': - if len(args) != 1: - print('ExecuteStatement requires 1 args') - sys.exit(1) - pp.pprint(client.ExecuteStatement(eval(args[0]),)) - -elif cmd == 'GetTypeInfo': - if len(args) != 1: - print('GetTypeInfo requires 1 args') - sys.exit(1) - pp.pprint(client.GetTypeInfo(eval(args[0]),)) - -elif cmd == 'GetCatalogs': - if len(args) != 1: - print('GetCatalogs requires 1 args') - sys.exit(1) - pp.pprint(client.GetCatalogs(eval(args[0]),)) - -elif cmd == 'GetSchemas': - if len(args) != 1: - print('GetSchemas requires 1 args') - sys.exit(1) - pp.pprint(client.GetSchemas(eval(args[0]),)) - -elif cmd == 'GetTables': - if len(args) != 1: - print('GetTables requires 1 args') - sys.exit(1) - pp.pprint(client.GetTables(eval(args[0]),)) - -elif cmd == 'GetTableTypes': - if len(args) != 1: - print('GetTableTypes requires 1 args') - sys.exit(1) - pp.pprint(client.GetTableTypes(eval(args[0]),)) - -elif cmd == 'GetColumns': - if len(args) != 1: - print('GetColumns requires 1 args') - sys.exit(1) - pp.pprint(client.GetColumns(eval(args[0]),)) - -elif cmd == 'GetFunctions': - if len(args) != 1: - print('GetFunctions requires 1 args') - sys.exit(1) - pp.pprint(client.GetFunctions(eval(args[0]),)) - -elif cmd == 'GetPrimaryKeys': - if len(args) != 1: - print('GetPrimaryKeys requires 1 args') - sys.exit(1) - pp.pprint(client.GetPrimaryKeys(eval(args[0]),)) - -elif cmd == 'GetCrossReference': - if len(args) != 1: - print('GetCrossReference requires 1 args') - sys.exit(1) - pp.pprint(client.GetCrossReference(eval(args[0]),)) - -elif cmd == 'GetOperationStatus': - if len(args) != 1: - print('GetOperationStatus requires 1 args') - sys.exit(1) - pp.pprint(client.GetOperationStatus(eval(args[0]),)) - -elif cmd == 'CancelOperation': - if len(args) != 1: - print('CancelOperation requires 1 args') - sys.exit(1) - pp.pprint(client.CancelOperation(eval(args[0]),)) - -elif cmd == 'CloseOperation': - if len(args) != 1: - print('CloseOperation requires 1 args') - sys.exit(1) - pp.pprint(client.CloseOperation(eval(args[0]),)) - -elif cmd == 'GetResultSetMetadata': - if len(args) != 1: - print('GetResultSetMetadata requires 1 args') - sys.exit(1) - pp.pprint(client.GetResultSetMetadata(eval(args[0]),)) - -elif cmd == 'FetchResults': - if len(args) != 1: - print('FetchResults requires 1 args') - sys.exit(1) - pp.pprint(client.FetchResults(eval(args[0]),)) - -elif cmd == 'GetDelegationToken': - if len(args) != 1: - print('GetDelegationToken requires 1 args') - sys.exit(1) - pp.pprint(client.GetDelegationToken(eval(args[0]),)) - -elif cmd == 'CancelDelegationToken': - if len(args) != 1: - print('CancelDelegationToken requires 1 args') - sys.exit(1) - pp.pprint(client.CancelDelegationToken(eval(args[0]),)) - -elif cmd == 'RenewDelegationToken': - if len(args) != 1: - print('RenewDelegationToken requires 1 args') - sys.exit(1) - pp.pprint(client.RenewDelegationToken(eval(args[0]),)) - -elif cmd == 'GetLog': - if len(args) != 1: - print('GetLog requires 1 args') - sys.exit(1) - pp.pprint(client.GetLog(eval(args[0]),)) - -else: - print('Unrecognized method %s' % cmd) - sys.exit(1) - -transport.close() diff --git a/src/chronify/_vendor/kyuubi/TCLIService/TCLIService.py b/src/chronify/_vendor/kyuubi/TCLIService/TCLIService.py deleted file mode 100644 index bb1415e..0000000 --- a/src/chronify/_vendor/kyuubi/TCLIService/TCLIService.py +++ /dev/null @@ -1,3986 +0,0 @@ -# -# Autogenerated by Thrift Compiler (0.10.0) -# -# DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING -# -# options string: py -# - -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException -from thrift.protocol.TProtocol import TProtocolException -import sys -import logging -from .ttypes import * -from thrift.Thrift import TProcessor -from thrift.transport import TTransport - - -class Iface(object): - def OpenSession(self, req): - """ - Parameters: - - req - """ - pass - - def CloseSession(self, req): - """ - Parameters: - - req - """ - pass - - def GetInfo(self, req): - """ - Parameters: - - req - """ - pass - - def ExecuteStatement(self, req): - """ - Parameters: - - req - """ - pass - - def GetTypeInfo(self, req): - """ - Parameters: - - req - """ - pass - - def GetCatalogs(self, req): - """ - Parameters: - - req - """ - pass - - def GetSchemas(self, req): - """ - Parameters: - - req - """ - pass - - def GetTables(self, req): - """ - Parameters: - - req - """ - pass - - def GetTableTypes(self, req): - """ - Parameters: - - req - """ - pass - - def GetColumns(self, req): - """ - Parameters: - - req - """ - pass - - def GetFunctions(self, req): - """ - Parameters: - - req - """ - pass - - def GetPrimaryKeys(self, req): - """ - Parameters: - - req - """ - pass - - def GetCrossReference(self, req): - """ - Parameters: - - req - """ - pass - - def GetOperationStatus(self, req): - """ - Parameters: - - req - """ - pass - - def CancelOperation(self, req): - """ - Parameters: - - req - """ - pass - - def CloseOperation(self, req): - """ - Parameters: - - req - """ - pass - - def GetResultSetMetadata(self, req): - """ - Parameters: - - req - """ - pass - - def FetchResults(self, req): - """ - Parameters: - - req - """ - pass - - def GetDelegationToken(self, req): - """ - Parameters: - - req - """ - pass - - def CancelDelegationToken(self, req): - """ - Parameters: - - req - """ - pass - - def RenewDelegationToken(self, req): - """ - Parameters: - - req - """ - pass - - def GetLog(self, req): - """ - Parameters: - - req - """ - pass - - -class Client(Iface): - def __init__(self, iprot, oprot=None): - self._iprot = self._oprot = iprot - if oprot is not None: - self._oprot = oprot - self._seqid = 0 - - def OpenSession(self, req): - """ - Parameters: - - req - """ - self.send_OpenSession(req) - return self.recv_OpenSession() - - def send_OpenSession(self, req): - self._oprot.writeMessageBegin('OpenSession', TMessageType.CALL, self._seqid) - args = OpenSession_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_OpenSession(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = OpenSession_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "OpenSession failed: unknown result") - - def CloseSession(self, req): - """ - Parameters: - - req - """ - self.send_CloseSession(req) - return self.recv_CloseSession() - - def send_CloseSession(self, req): - self._oprot.writeMessageBegin('CloseSession', TMessageType.CALL, self._seqid) - args = CloseSession_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_CloseSession(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = CloseSession_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "CloseSession failed: unknown result") - - def GetInfo(self, req): - """ - Parameters: - - req - """ - self.send_GetInfo(req) - return self.recv_GetInfo() - - def send_GetInfo(self, req): - self._oprot.writeMessageBegin('GetInfo', TMessageType.CALL, self._seqid) - args = GetInfo_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetInfo(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetInfo_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetInfo failed: unknown result") - - def ExecuteStatement(self, req): - """ - Parameters: - - req - """ - self.send_ExecuteStatement(req) - return self.recv_ExecuteStatement() - - def send_ExecuteStatement(self, req): - self._oprot.writeMessageBegin('ExecuteStatement', TMessageType.CALL, self._seqid) - args = ExecuteStatement_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_ExecuteStatement(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = ExecuteStatement_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "ExecuteStatement failed: unknown result") - - def GetTypeInfo(self, req): - """ - Parameters: - - req - """ - self.send_GetTypeInfo(req) - return self.recv_GetTypeInfo() - - def send_GetTypeInfo(self, req): - self._oprot.writeMessageBegin('GetTypeInfo', TMessageType.CALL, self._seqid) - args = GetTypeInfo_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetTypeInfo(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetTypeInfo_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetTypeInfo failed: unknown result") - - def GetCatalogs(self, req): - """ - Parameters: - - req - """ - self.send_GetCatalogs(req) - return self.recv_GetCatalogs() - - def send_GetCatalogs(self, req): - self._oprot.writeMessageBegin('GetCatalogs', TMessageType.CALL, self._seqid) - args = GetCatalogs_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetCatalogs(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetCatalogs_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetCatalogs failed: unknown result") - - def GetSchemas(self, req): - """ - Parameters: - - req - """ - self.send_GetSchemas(req) - return self.recv_GetSchemas() - - def send_GetSchemas(self, req): - self._oprot.writeMessageBegin('GetSchemas', TMessageType.CALL, self._seqid) - args = GetSchemas_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetSchemas(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetSchemas_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetSchemas failed: unknown result") - - def GetTables(self, req): - """ - Parameters: - - req - """ - self.send_GetTables(req) - return self.recv_GetTables() - - def send_GetTables(self, req): - self._oprot.writeMessageBegin('GetTables', TMessageType.CALL, self._seqid) - args = GetTables_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetTables(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetTables_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetTables failed: unknown result") - - def GetTableTypes(self, req): - """ - Parameters: - - req - """ - self.send_GetTableTypes(req) - return self.recv_GetTableTypes() - - def send_GetTableTypes(self, req): - self._oprot.writeMessageBegin('GetTableTypes', TMessageType.CALL, self._seqid) - args = GetTableTypes_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetTableTypes(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetTableTypes_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetTableTypes failed: unknown result") - - def GetColumns(self, req): - """ - Parameters: - - req - """ - self.send_GetColumns(req) - return self.recv_GetColumns() - - def send_GetColumns(self, req): - self._oprot.writeMessageBegin('GetColumns', TMessageType.CALL, self._seqid) - args = GetColumns_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetColumns(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetColumns_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetColumns failed: unknown result") - - def GetFunctions(self, req): - """ - Parameters: - - req - """ - self.send_GetFunctions(req) - return self.recv_GetFunctions() - - def send_GetFunctions(self, req): - self._oprot.writeMessageBegin('GetFunctions', TMessageType.CALL, self._seqid) - args = GetFunctions_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetFunctions(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetFunctions_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetFunctions failed: unknown result") - - def GetPrimaryKeys(self, req): - """ - Parameters: - - req - """ - self.send_GetPrimaryKeys(req) - return self.recv_GetPrimaryKeys() - - def send_GetPrimaryKeys(self, req): - self._oprot.writeMessageBegin('GetPrimaryKeys', TMessageType.CALL, self._seqid) - args = GetPrimaryKeys_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetPrimaryKeys(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetPrimaryKeys_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetPrimaryKeys failed: unknown result") - - def GetCrossReference(self, req): - """ - Parameters: - - req - """ - self.send_GetCrossReference(req) - return self.recv_GetCrossReference() - - def send_GetCrossReference(self, req): - self._oprot.writeMessageBegin('GetCrossReference', TMessageType.CALL, self._seqid) - args = GetCrossReference_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetCrossReference(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetCrossReference_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetCrossReference failed: unknown result") - - def GetOperationStatus(self, req): - """ - Parameters: - - req - """ - self.send_GetOperationStatus(req) - return self.recv_GetOperationStatus() - - def send_GetOperationStatus(self, req): - self._oprot.writeMessageBegin('GetOperationStatus', TMessageType.CALL, self._seqid) - args = GetOperationStatus_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetOperationStatus(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetOperationStatus_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetOperationStatus failed: unknown result") - - def CancelOperation(self, req): - """ - Parameters: - - req - """ - self.send_CancelOperation(req) - return self.recv_CancelOperation() - - def send_CancelOperation(self, req): - self._oprot.writeMessageBegin('CancelOperation', TMessageType.CALL, self._seqid) - args = CancelOperation_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_CancelOperation(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = CancelOperation_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "CancelOperation failed: unknown result") - - def CloseOperation(self, req): - """ - Parameters: - - req - """ - self.send_CloseOperation(req) - return self.recv_CloseOperation() - - def send_CloseOperation(self, req): - self._oprot.writeMessageBegin('CloseOperation', TMessageType.CALL, self._seqid) - args = CloseOperation_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_CloseOperation(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = CloseOperation_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "CloseOperation failed: unknown result") - - def GetResultSetMetadata(self, req): - """ - Parameters: - - req - """ - self.send_GetResultSetMetadata(req) - return self.recv_GetResultSetMetadata() - - def send_GetResultSetMetadata(self, req): - self._oprot.writeMessageBegin('GetResultSetMetadata', TMessageType.CALL, self._seqid) - args = GetResultSetMetadata_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetResultSetMetadata(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetResultSetMetadata_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetResultSetMetadata failed: unknown result") - - def FetchResults(self, req): - """ - Parameters: - - req - """ - self.send_FetchResults(req) - return self.recv_FetchResults() - - def send_FetchResults(self, req): - self._oprot.writeMessageBegin('FetchResults', TMessageType.CALL, self._seqid) - args = FetchResults_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_FetchResults(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = FetchResults_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "FetchResults failed: unknown result") - - def GetDelegationToken(self, req): - """ - Parameters: - - req - """ - self.send_GetDelegationToken(req) - return self.recv_GetDelegationToken() - - def send_GetDelegationToken(self, req): - self._oprot.writeMessageBegin('GetDelegationToken', TMessageType.CALL, self._seqid) - args = GetDelegationToken_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetDelegationToken(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetDelegationToken_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetDelegationToken failed: unknown result") - - def CancelDelegationToken(self, req): - """ - Parameters: - - req - """ - self.send_CancelDelegationToken(req) - return self.recv_CancelDelegationToken() - - def send_CancelDelegationToken(self, req): - self._oprot.writeMessageBegin('CancelDelegationToken', TMessageType.CALL, self._seqid) - args = CancelDelegationToken_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_CancelDelegationToken(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = CancelDelegationToken_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "CancelDelegationToken failed: unknown result") - - def RenewDelegationToken(self, req): - """ - Parameters: - - req - """ - self.send_RenewDelegationToken(req) - return self.recv_RenewDelegationToken() - - def send_RenewDelegationToken(self, req): - self._oprot.writeMessageBegin('RenewDelegationToken', TMessageType.CALL, self._seqid) - args = RenewDelegationToken_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_RenewDelegationToken(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = RenewDelegationToken_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "RenewDelegationToken failed: unknown result") - - def GetLog(self, req): - """ - Parameters: - - req - """ - self.send_GetLog(req) - return self.recv_GetLog() - - def send_GetLog(self, req): - self._oprot.writeMessageBegin('GetLog', TMessageType.CALL, self._seqid) - args = GetLog_args() - args.req = req - args.write(self._oprot) - self._oprot.writeMessageEnd() - self._oprot.trans.flush() - - def recv_GetLog(self): - iprot = self._iprot - (fname, mtype, rseqid) = iprot.readMessageBegin() - if mtype == TMessageType.EXCEPTION: - x = TApplicationException() - x.read(iprot) - iprot.readMessageEnd() - raise x - result = GetLog_result() - result.read(iprot) - iprot.readMessageEnd() - if result.success is not None: - return result.success - raise TApplicationException(TApplicationException.MISSING_RESULT, "GetLog failed: unknown result") - - -class Processor(Iface, TProcessor): - def __init__(self, handler): - self._handler = handler - self._processMap = {} - self._processMap["OpenSession"] = Processor.process_OpenSession - self._processMap["CloseSession"] = Processor.process_CloseSession - self._processMap["GetInfo"] = Processor.process_GetInfo - self._processMap["ExecuteStatement"] = Processor.process_ExecuteStatement - self._processMap["GetTypeInfo"] = Processor.process_GetTypeInfo - self._processMap["GetCatalogs"] = Processor.process_GetCatalogs - self._processMap["GetSchemas"] = Processor.process_GetSchemas - self._processMap["GetTables"] = Processor.process_GetTables - self._processMap["GetTableTypes"] = Processor.process_GetTableTypes - self._processMap["GetColumns"] = Processor.process_GetColumns - self._processMap["GetFunctions"] = Processor.process_GetFunctions - self._processMap["GetPrimaryKeys"] = Processor.process_GetPrimaryKeys - self._processMap["GetCrossReference"] = Processor.process_GetCrossReference - self._processMap["GetOperationStatus"] = Processor.process_GetOperationStatus - self._processMap["CancelOperation"] = Processor.process_CancelOperation - self._processMap["CloseOperation"] = Processor.process_CloseOperation - self._processMap["GetResultSetMetadata"] = Processor.process_GetResultSetMetadata - self._processMap["FetchResults"] = Processor.process_FetchResults - self._processMap["GetDelegationToken"] = Processor.process_GetDelegationToken - self._processMap["CancelDelegationToken"] = Processor.process_CancelDelegationToken - self._processMap["RenewDelegationToken"] = Processor.process_RenewDelegationToken - self._processMap["GetLog"] = Processor.process_GetLog - - def process(self, iprot, oprot): - (name, type, seqid) = iprot.readMessageBegin() - if name not in self._processMap: - iprot.skip(TType.STRUCT) - iprot.readMessageEnd() - x = TApplicationException(TApplicationException.UNKNOWN_METHOD, 'Unknown function %s' % (name)) - oprot.writeMessageBegin(name, TMessageType.EXCEPTION, seqid) - x.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - return - else: - self._processMap[name](self, seqid, iprot, oprot) - return True - - def process_OpenSession(self, seqid, iprot, oprot): - args = OpenSession_args() - args.read(iprot) - iprot.readMessageEnd() - result = OpenSession_result() - try: - result.success = self._handler.OpenSession(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("OpenSession", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_CloseSession(self, seqid, iprot, oprot): - args = CloseSession_args() - args.read(iprot) - iprot.readMessageEnd() - result = CloseSession_result() - try: - result.success = self._handler.CloseSession(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("CloseSession", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetInfo(self, seqid, iprot, oprot): - args = GetInfo_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetInfo_result() - try: - result.success = self._handler.GetInfo(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetInfo", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_ExecuteStatement(self, seqid, iprot, oprot): - args = ExecuteStatement_args() - args.read(iprot) - iprot.readMessageEnd() - result = ExecuteStatement_result() - try: - result.success = self._handler.ExecuteStatement(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("ExecuteStatement", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetTypeInfo(self, seqid, iprot, oprot): - args = GetTypeInfo_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetTypeInfo_result() - try: - result.success = self._handler.GetTypeInfo(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetTypeInfo", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetCatalogs(self, seqid, iprot, oprot): - args = GetCatalogs_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetCatalogs_result() - try: - result.success = self._handler.GetCatalogs(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetCatalogs", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetSchemas(self, seqid, iprot, oprot): - args = GetSchemas_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetSchemas_result() - try: - result.success = self._handler.GetSchemas(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetSchemas", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetTables(self, seqid, iprot, oprot): - args = GetTables_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetTables_result() - try: - result.success = self._handler.GetTables(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetTables", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetTableTypes(self, seqid, iprot, oprot): - args = GetTableTypes_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetTableTypes_result() - try: - result.success = self._handler.GetTableTypes(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetTableTypes", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetColumns(self, seqid, iprot, oprot): - args = GetColumns_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetColumns_result() - try: - result.success = self._handler.GetColumns(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetColumns", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetFunctions(self, seqid, iprot, oprot): - args = GetFunctions_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetFunctions_result() - try: - result.success = self._handler.GetFunctions(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetFunctions", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetPrimaryKeys(self, seqid, iprot, oprot): - args = GetPrimaryKeys_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetPrimaryKeys_result() - try: - result.success = self._handler.GetPrimaryKeys(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetPrimaryKeys", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetCrossReference(self, seqid, iprot, oprot): - args = GetCrossReference_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetCrossReference_result() - try: - result.success = self._handler.GetCrossReference(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetCrossReference", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetOperationStatus(self, seqid, iprot, oprot): - args = GetOperationStatus_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetOperationStatus_result() - try: - result.success = self._handler.GetOperationStatus(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetOperationStatus", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_CancelOperation(self, seqid, iprot, oprot): - args = CancelOperation_args() - args.read(iprot) - iprot.readMessageEnd() - result = CancelOperation_result() - try: - result.success = self._handler.CancelOperation(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("CancelOperation", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_CloseOperation(self, seqid, iprot, oprot): - args = CloseOperation_args() - args.read(iprot) - iprot.readMessageEnd() - result = CloseOperation_result() - try: - result.success = self._handler.CloseOperation(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("CloseOperation", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetResultSetMetadata(self, seqid, iprot, oprot): - args = GetResultSetMetadata_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetResultSetMetadata_result() - try: - result.success = self._handler.GetResultSetMetadata(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetResultSetMetadata", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_FetchResults(self, seqid, iprot, oprot): - args = FetchResults_args() - args.read(iprot) - iprot.readMessageEnd() - result = FetchResults_result() - try: - result.success = self._handler.FetchResults(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("FetchResults", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetDelegationToken(self, seqid, iprot, oprot): - args = GetDelegationToken_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetDelegationToken_result() - try: - result.success = self._handler.GetDelegationToken(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetDelegationToken", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_CancelDelegationToken(self, seqid, iprot, oprot): - args = CancelDelegationToken_args() - args.read(iprot) - iprot.readMessageEnd() - result = CancelDelegationToken_result() - try: - result.success = self._handler.CancelDelegationToken(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("CancelDelegationToken", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_RenewDelegationToken(self, seqid, iprot, oprot): - args = RenewDelegationToken_args() - args.read(iprot) - iprot.readMessageEnd() - result = RenewDelegationToken_result() - try: - result.success = self._handler.RenewDelegationToken(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("RenewDelegationToken", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - - def process_GetLog(self, seqid, iprot, oprot): - args = GetLog_args() - args.read(iprot) - iprot.readMessageEnd() - result = GetLog_result() - try: - result.success = self._handler.GetLog(args.req) - msg_type = TMessageType.REPLY - except (TTransport.TTransportException, KeyboardInterrupt, SystemExit): - raise - except Exception as ex: - msg_type = TMessageType.EXCEPTION - logging.exception(ex) - result = TApplicationException(TApplicationException.INTERNAL_ERROR, 'Internal error') - oprot.writeMessageBegin("GetLog", msg_type, seqid) - result.write(oprot) - oprot.writeMessageEnd() - oprot.trans.flush() - -# HELPER FUNCTIONS AND STRUCTURES - - -class OpenSession_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TOpenSessionReq, TOpenSessionReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TOpenSessionReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('OpenSession_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class OpenSession_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TOpenSessionResp, TOpenSessionResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TOpenSessionResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('OpenSession_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class CloseSession_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TCloseSessionReq, TCloseSessionReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TCloseSessionReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('CloseSession_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class CloseSession_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TCloseSessionResp, TCloseSessionResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TCloseSessionResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('CloseSession_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetInfo_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetInfoReq, TGetInfoReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetInfoReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetInfo_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetInfo_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetInfoResp, TGetInfoResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetInfoResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetInfo_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class ExecuteStatement_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TExecuteStatementReq, TExecuteStatementReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TExecuteStatementReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('ExecuteStatement_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class ExecuteStatement_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TExecuteStatementResp, TExecuteStatementResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TExecuteStatementResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('ExecuteStatement_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetTypeInfo_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetTypeInfoReq, TGetTypeInfoReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetTypeInfoReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetTypeInfo_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetTypeInfo_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetTypeInfoResp, TGetTypeInfoResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetTypeInfoResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetTypeInfo_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetCatalogs_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetCatalogsReq, TGetCatalogsReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetCatalogsReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetCatalogs_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetCatalogs_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetCatalogsResp, TGetCatalogsResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetCatalogsResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetCatalogs_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetSchemas_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetSchemasReq, TGetSchemasReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetSchemasReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetSchemas_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetSchemas_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetSchemasResp, TGetSchemasResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetSchemasResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetSchemas_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetTables_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetTablesReq, TGetTablesReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetTablesReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetTables_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetTables_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetTablesResp, TGetTablesResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetTablesResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetTables_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetTableTypes_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetTableTypesReq, TGetTableTypesReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetTableTypesReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetTableTypes_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetTableTypes_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetTableTypesResp, TGetTableTypesResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetTableTypesResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetTableTypes_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetColumns_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetColumnsReq, TGetColumnsReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetColumnsReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetColumns_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetColumns_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetColumnsResp, TGetColumnsResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetColumnsResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetColumns_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetFunctions_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetFunctionsReq, TGetFunctionsReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetFunctionsReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetFunctions_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetFunctions_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetFunctionsResp, TGetFunctionsResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetFunctionsResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetFunctions_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetPrimaryKeys_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetPrimaryKeysReq, TGetPrimaryKeysReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetPrimaryKeysReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetPrimaryKeys_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetPrimaryKeys_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetPrimaryKeysResp, TGetPrimaryKeysResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetPrimaryKeysResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetPrimaryKeys_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetCrossReference_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetCrossReferenceReq, TGetCrossReferenceReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetCrossReferenceReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetCrossReference_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetCrossReference_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetCrossReferenceResp, TGetCrossReferenceResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetCrossReferenceResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetCrossReference_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetOperationStatus_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetOperationStatusReq, TGetOperationStatusReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetOperationStatusReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetOperationStatus_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetOperationStatus_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetOperationStatusResp, TGetOperationStatusResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetOperationStatusResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetOperationStatus_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class CancelOperation_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TCancelOperationReq, TCancelOperationReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TCancelOperationReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('CancelOperation_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class CancelOperation_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TCancelOperationResp, TCancelOperationResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TCancelOperationResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('CancelOperation_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class CloseOperation_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TCloseOperationReq, TCloseOperationReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TCloseOperationReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('CloseOperation_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class CloseOperation_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TCloseOperationResp, TCloseOperationResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TCloseOperationResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('CloseOperation_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetResultSetMetadata_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetResultSetMetadataReq, TGetResultSetMetadataReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetResultSetMetadataReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetResultSetMetadata_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetResultSetMetadata_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetResultSetMetadataResp, TGetResultSetMetadataResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetResultSetMetadataResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetResultSetMetadata_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class FetchResults_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TFetchResultsReq, TFetchResultsReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TFetchResultsReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('FetchResults_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class FetchResults_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TFetchResultsResp, TFetchResultsResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TFetchResultsResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('FetchResults_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetDelegationToken_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetDelegationTokenReq, TGetDelegationTokenReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetDelegationTokenReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetDelegationToken_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetDelegationToken_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetDelegationTokenResp, TGetDelegationTokenResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetDelegationTokenResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetDelegationToken_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class CancelDelegationToken_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TCancelDelegationTokenReq, TCancelDelegationTokenReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TCancelDelegationTokenReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('CancelDelegationToken_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class CancelDelegationToken_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TCancelDelegationTokenResp, TCancelDelegationTokenResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TCancelDelegationTokenResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('CancelDelegationToken_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class RenewDelegationToken_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TRenewDelegationTokenReq, TRenewDelegationTokenReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TRenewDelegationTokenReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('RenewDelegationToken_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class RenewDelegationToken_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TRenewDelegationTokenResp, TRenewDelegationTokenResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TRenewDelegationTokenResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('RenewDelegationToken_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetLog_args(object): - """ - Attributes: - - req - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'req', (TGetLogReq, TGetLogReq.thrift_spec), None, ), # 1 - ) - - def __init__(self, req=None,): - self.req = req - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.req = TGetLogReq() - self.req.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetLog_args') - if self.req is not None: - oprot.writeFieldBegin('req', TType.STRUCT, 1) - self.req.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class GetLog_result(object): - """ - Attributes: - - success - """ - - thrift_spec = ( - (0, TType.STRUCT, 'success', (TGetLogResp, TGetLogResp.thrift_spec), None, ), # 0 - ) - - def __init__(self, success=None,): - self.success = success - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 0: - if ftype == TType.STRUCT: - self.success = TGetLogResp() - self.success.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('GetLog_result') - if self.success is not None: - oprot.writeFieldBegin('success', TType.STRUCT, 0) - self.success.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) diff --git a/src/chronify/_vendor/kyuubi/TCLIService/__init__.py b/src/chronify/_vendor/kyuubi/TCLIService/__init__.py deleted file mode 100644 index fe5e7c2..0000000 --- a/src/chronify/_vendor/kyuubi/TCLIService/__init__.py +++ /dev/null @@ -1 +0,0 @@ -__all__ = ['ttypes', 'constants', 'TCLIService'] diff --git a/src/chronify/_vendor/kyuubi/TCLIService/constants.py b/src/chronify/_vendor/kyuubi/TCLIService/constants.py deleted file mode 100644 index b6bf88e..0000000 --- a/src/chronify/_vendor/kyuubi/TCLIService/constants.py +++ /dev/null @@ -1,68 +0,0 @@ -# -# Autogenerated by Thrift Compiler (0.10.0) -# -# DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING -# -# options string: py -# - -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException -from thrift.protocol.TProtocol import TProtocolException -import sys -from .ttypes import * -PRIMITIVE_TYPES = set(( - 0, - 1, - 2, - 3, - 4, - 5, - 6, - 7, - 8, - 9, - 15, - 16, - 17, - 18, - 19, - 20, - 21, -)) -COMPLEX_TYPES = set(( - 10, - 11, - 12, - 13, - 14, -)) -COLLECTION_TYPES = set(( - 10, - 11, -)) -TYPE_NAMES = { - 0: "BOOLEAN", - 1: "TINYINT", - 2: "SMALLINT", - 3: "INT", - 4: "BIGINT", - 5: "FLOAT", - 6: "DOUBLE", - 7: "STRING", - 8: "TIMESTAMP", - 9: "BINARY", - 10: "ARRAY", - 11: "MAP", - 12: "STRUCT", - 13: "UNIONTYPE", - 15: "DECIMAL", - 16: "NULL", - 17: "DATE", - 18: "VARCHAR", - 19: "CHAR", - 20: "INTERVAL_YEAR_MONTH", - 21: "INTERVAL_DAY_TIME", -} -CHARACTER_MAXIMUM_LENGTH = "characterMaximumLength" -PRECISION = "precision" -SCALE = "scale" diff --git a/src/chronify/_vendor/kyuubi/TCLIService/ttypes.py b/src/chronify/_vendor/kyuubi/TCLIService/ttypes.py deleted file mode 100644 index 573bd04..0000000 --- a/src/chronify/_vendor/kyuubi/TCLIService/ttypes.py +++ /dev/null @@ -1,7210 +0,0 @@ -# -# Autogenerated by Thrift Compiler (0.10.0) -# -# DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING -# -# options string: py -# - -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException -from thrift.protocol.TProtocol import TProtocolException -import sys - -from thrift.transport import TTransport - - -class TProtocolVersion(object): - HIVE_CLI_SERVICE_PROTOCOL_V1 = 0 - HIVE_CLI_SERVICE_PROTOCOL_V2 = 1 - HIVE_CLI_SERVICE_PROTOCOL_V3 = 2 - HIVE_CLI_SERVICE_PROTOCOL_V4 = 3 - HIVE_CLI_SERVICE_PROTOCOL_V5 = 4 - HIVE_CLI_SERVICE_PROTOCOL_V6 = 5 - HIVE_CLI_SERVICE_PROTOCOL_V7 = 6 - HIVE_CLI_SERVICE_PROTOCOL_V8 = 7 - HIVE_CLI_SERVICE_PROTOCOL_V9 = 8 - HIVE_CLI_SERVICE_PROTOCOL_V10 = 9 - - _VALUES_TO_NAMES = { - 0: "HIVE_CLI_SERVICE_PROTOCOL_V1", - 1: "HIVE_CLI_SERVICE_PROTOCOL_V2", - 2: "HIVE_CLI_SERVICE_PROTOCOL_V3", - 3: "HIVE_CLI_SERVICE_PROTOCOL_V4", - 4: "HIVE_CLI_SERVICE_PROTOCOL_V5", - 5: "HIVE_CLI_SERVICE_PROTOCOL_V6", - 6: "HIVE_CLI_SERVICE_PROTOCOL_V7", - 7: "HIVE_CLI_SERVICE_PROTOCOL_V8", - 8: "HIVE_CLI_SERVICE_PROTOCOL_V9", - 9: "HIVE_CLI_SERVICE_PROTOCOL_V10", - } - - _NAMES_TO_VALUES = { - "HIVE_CLI_SERVICE_PROTOCOL_V1": 0, - "HIVE_CLI_SERVICE_PROTOCOL_V2": 1, - "HIVE_CLI_SERVICE_PROTOCOL_V3": 2, - "HIVE_CLI_SERVICE_PROTOCOL_V4": 3, - "HIVE_CLI_SERVICE_PROTOCOL_V5": 4, - "HIVE_CLI_SERVICE_PROTOCOL_V6": 5, - "HIVE_CLI_SERVICE_PROTOCOL_V7": 6, - "HIVE_CLI_SERVICE_PROTOCOL_V8": 7, - "HIVE_CLI_SERVICE_PROTOCOL_V9": 8, - "HIVE_CLI_SERVICE_PROTOCOL_V10": 9, - } - - -class TTypeId(object): - BOOLEAN_TYPE = 0 - TINYINT_TYPE = 1 - SMALLINT_TYPE = 2 - INT_TYPE = 3 - BIGINT_TYPE = 4 - FLOAT_TYPE = 5 - DOUBLE_TYPE = 6 - STRING_TYPE = 7 - TIMESTAMP_TYPE = 8 - BINARY_TYPE = 9 - ARRAY_TYPE = 10 - MAP_TYPE = 11 - STRUCT_TYPE = 12 - UNION_TYPE = 13 - USER_DEFINED_TYPE = 14 - DECIMAL_TYPE = 15 - NULL_TYPE = 16 - DATE_TYPE = 17 - VARCHAR_TYPE = 18 - CHAR_TYPE = 19 - INTERVAL_YEAR_MONTH_TYPE = 20 - INTERVAL_DAY_TIME_TYPE = 21 - - _VALUES_TO_NAMES = { - 0: "BOOLEAN_TYPE", - 1: "TINYINT_TYPE", - 2: "SMALLINT_TYPE", - 3: "INT_TYPE", - 4: "BIGINT_TYPE", - 5: "FLOAT_TYPE", - 6: "DOUBLE_TYPE", - 7: "STRING_TYPE", - 8: "TIMESTAMP_TYPE", - 9: "BINARY_TYPE", - 10: "ARRAY_TYPE", - 11: "MAP_TYPE", - 12: "STRUCT_TYPE", - 13: "UNION_TYPE", - 14: "USER_DEFINED_TYPE", - 15: "DECIMAL_TYPE", - 16: "NULL_TYPE", - 17: "DATE_TYPE", - 18: "VARCHAR_TYPE", - 19: "CHAR_TYPE", - 20: "INTERVAL_YEAR_MONTH_TYPE", - 21: "INTERVAL_DAY_TIME_TYPE", - } - - _NAMES_TO_VALUES = { - "BOOLEAN_TYPE": 0, - "TINYINT_TYPE": 1, - "SMALLINT_TYPE": 2, - "INT_TYPE": 3, - "BIGINT_TYPE": 4, - "FLOAT_TYPE": 5, - "DOUBLE_TYPE": 6, - "STRING_TYPE": 7, - "TIMESTAMP_TYPE": 8, - "BINARY_TYPE": 9, - "ARRAY_TYPE": 10, - "MAP_TYPE": 11, - "STRUCT_TYPE": 12, - "UNION_TYPE": 13, - "USER_DEFINED_TYPE": 14, - "DECIMAL_TYPE": 15, - "NULL_TYPE": 16, - "DATE_TYPE": 17, - "VARCHAR_TYPE": 18, - "CHAR_TYPE": 19, - "INTERVAL_YEAR_MONTH_TYPE": 20, - "INTERVAL_DAY_TIME_TYPE": 21, - } - - -class TStatusCode(object): - SUCCESS_STATUS = 0 - SUCCESS_WITH_INFO_STATUS = 1 - STILL_EXECUTING_STATUS = 2 - ERROR_STATUS = 3 - INVALID_HANDLE_STATUS = 4 - - _VALUES_TO_NAMES = { - 0: "SUCCESS_STATUS", - 1: "SUCCESS_WITH_INFO_STATUS", - 2: "STILL_EXECUTING_STATUS", - 3: "ERROR_STATUS", - 4: "INVALID_HANDLE_STATUS", - } - - _NAMES_TO_VALUES = { - "SUCCESS_STATUS": 0, - "SUCCESS_WITH_INFO_STATUS": 1, - "STILL_EXECUTING_STATUS": 2, - "ERROR_STATUS": 3, - "INVALID_HANDLE_STATUS": 4, - } - - -class TOperationState(object): - INITIALIZED_STATE = 0 - RUNNING_STATE = 1 - FINISHED_STATE = 2 - CANCELED_STATE = 3 - CLOSED_STATE = 4 - ERROR_STATE = 5 - UKNOWN_STATE = 6 - PENDING_STATE = 7 - TIMEDOUT_STATE = 8 - - _VALUES_TO_NAMES = { - 0: "INITIALIZED_STATE", - 1: "RUNNING_STATE", - 2: "FINISHED_STATE", - 3: "CANCELED_STATE", - 4: "CLOSED_STATE", - 5: "ERROR_STATE", - 6: "UKNOWN_STATE", - 7: "PENDING_STATE", - 8: "TIMEDOUT_STATE", - } - - _NAMES_TO_VALUES = { - "INITIALIZED_STATE": 0, - "RUNNING_STATE": 1, - "FINISHED_STATE": 2, - "CANCELED_STATE": 3, - "CLOSED_STATE": 4, - "ERROR_STATE": 5, - "UKNOWN_STATE": 6, - "PENDING_STATE": 7, - "TIMEDOUT_STATE": 8, - } - - -class TOperationType(object): - EXECUTE_STATEMENT = 0 - GET_TYPE_INFO = 1 - GET_CATALOGS = 2 - GET_SCHEMAS = 3 - GET_TABLES = 4 - GET_TABLE_TYPES = 5 - GET_COLUMNS = 6 - GET_FUNCTIONS = 7 - UNKNOWN = 8 - - _VALUES_TO_NAMES = { - 0: "EXECUTE_STATEMENT", - 1: "GET_TYPE_INFO", - 2: "GET_CATALOGS", - 3: "GET_SCHEMAS", - 4: "GET_TABLES", - 5: "GET_TABLE_TYPES", - 6: "GET_COLUMNS", - 7: "GET_FUNCTIONS", - 8: "UNKNOWN", - } - - _NAMES_TO_VALUES = { - "EXECUTE_STATEMENT": 0, - "GET_TYPE_INFO": 1, - "GET_CATALOGS": 2, - "GET_SCHEMAS": 3, - "GET_TABLES": 4, - "GET_TABLE_TYPES": 5, - "GET_COLUMNS": 6, - "GET_FUNCTIONS": 7, - "UNKNOWN": 8, - } - - -class TGetInfoType(object): - CLI_MAX_DRIVER_CONNECTIONS = 0 - CLI_MAX_CONCURRENT_ACTIVITIES = 1 - CLI_DATA_SOURCE_NAME = 2 - CLI_FETCH_DIRECTION = 8 - CLI_SERVER_NAME = 13 - CLI_SEARCH_PATTERN_ESCAPE = 14 - CLI_DBMS_NAME = 17 - CLI_DBMS_VER = 18 - CLI_ACCESSIBLE_TABLES = 19 - CLI_ACCESSIBLE_PROCEDURES = 20 - CLI_CURSOR_COMMIT_BEHAVIOR = 23 - CLI_DATA_SOURCE_READ_ONLY = 25 - CLI_DEFAULT_TXN_ISOLATION = 26 - CLI_IDENTIFIER_CASE = 28 - CLI_IDENTIFIER_QUOTE_CHAR = 29 - CLI_MAX_COLUMN_NAME_LEN = 30 - CLI_MAX_CURSOR_NAME_LEN = 31 - CLI_MAX_SCHEMA_NAME_LEN = 32 - CLI_MAX_CATALOG_NAME_LEN = 34 - CLI_MAX_TABLE_NAME_LEN = 35 - CLI_SCROLL_CONCURRENCY = 43 - CLI_TXN_CAPABLE = 46 - CLI_USER_NAME = 47 - CLI_TXN_ISOLATION_OPTION = 72 - CLI_INTEGRITY = 73 - CLI_GETDATA_EXTENSIONS = 81 - CLI_NULL_COLLATION = 85 - CLI_ALTER_TABLE = 86 - CLI_ORDER_BY_COLUMNS_IN_SELECT = 90 - CLI_SPECIAL_CHARACTERS = 94 - CLI_MAX_COLUMNS_IN_GROUP_BY = 97 - CLI_MAX_COLUMNS_IN_INDEX = 98 - CLI_MAX_COLUMNS_IN_ORDER_BY = 99 - CLI_MAX_COLUMNS_IN_SELECT = 100 - CLI_MAX_COLUMNS_IN_TABLE = 101 - CLI_MAX_INDEX_SIZE = 102 - CLI_MAX_ROW_SIZE = 104 - CLI_MAX_STATEMENT_LEN = 105 - CLI_MAX_TABLES_IN_SELECT = 106 - CLI_MAX_USER_NAME_LEN = 107 - CLI_OJ_CAPABILITIES = 115 - CLI_XOPEN_CLI_YEAR = 10000 - CLI_CURSOR_SENSITIVITY = 10001 - CLI_DESCRIBE_PARAMETER = 10002 - CLI_CATALOG_NAME = 10003 - CLI_COLLATION_SEQ = 10004 - CLI_MAX_IDENTIFIER_LEN = 10005 - - _VALUES_TO_NAMES = { - 0: "CLI_MAX_DRIVER_CONNECTIONS", - 1: "CLI_MAX_CONCURRENT_ACTIVITIES", - 2: "CLI_DATA_SOURCE_NAME", - 8: "CLI_FETCH_DIRECTION", - 13: "CLI_SERVER_NAME", - 14: "CLI_SEARCH_PATTERN_ESCAPE", - 17: "CLI_DBMS_NAME", - 18: "CLI_DBMS_VER", - 19: "CLI_ACCESSIBLE_TABLES", - 20: "CLI_ACCESSIBLE_PROCEDURES", - 23: "CLI_CURSOR_COMMIT_BEHAVIOR", - 25: "CLI_DATA_SOURCE_READ_ONLY", - 26: "CLI_DEFAULT_TXN_ISOLATION", - 28: "CLI_IDENTIFIER_CASE", - 29: "CLI_IDENTIFIER_QUOTE_CHAR", - 30: "CLI_MAX_COLUMN_NAME_LEN", - 31: "CLI_MAX_CURSOR_NAME_LEN", - 32: "CLI_MAX_SCHEMA_NAME_LEN", - 34: "CLI_MAX_CATALOG_NAME_LEN", - 35: "CLI_MAX_TABLE_NAME_LEN", - 43: "CLI_SCROLL_CONCURRENCY", - 46: "CLI_TXN_CAPABLE", - 47: "CLI_USER_NAME", - 72: "CLI_TXN_ISOLATION_OPTION", - 73: "CLI_INTEGRITY", - 81: "CLI_GETDATA_EXTENSIONS", - 85: "CLI_NULL_COLLATION", - 86: "CLI_ALTER_TABLE", - 90: "CLI_ORDER_BY_COLUMNS_IN_SELECT", - 94: "CLI_SPECIAL_CHARACTERS", - 97: "CLI_MAX_COLUMNS_IN_GROUP_BY", - 98: "CLI_MAX_COLUMNS_IN_INDEX", - 99: "CLI_MAX_COLUMNS_IN_ORDER_BY", - 100: "CLI_MAX_COLUMNS_IN_SELECT", - 101: "CLI_MAX_COLUMNS_IN_TABLE", - 102: "CLI_MAX_INDEX_SIZE", - 104: "CLI_MAX_ROW_SIZE", - 105: "CLI_MAX_STATEMENT_LEN", - 106: "CLI_MAX_TABLES_IN_SELECT", - 107: "CLI_MAX_USER_NAME_LEN", - 115: "CLI_OJ_CAPABILITIES", - 10000: "CLI_XOPEN_CLI_YEAR", - 10001: "CLI_CURSOR_SENSITIVITY", - 10002: "CLI_DESCRIBE_PARAMETER", - 10003: "CLI_CATALOG_NAME", - 10004: "CLI_COLLATION_SEQ", - 10005: "CLI_MAX_IDENTIFIER_LEN", - } - - _NAMES_TO_VALUES = { - "CLI_MAX_DRIVER_CONNECTIONS": 0, - "CLI_MAX_CONCURRENT_ACTIVITIES": 1, - "CLI_DATA_SOURCE_NAME": 2, - "CLI_FETCH_DIRECTION": 8, - "CLI_SERVER_NAME": 13, - "CLI_SEARCH_PATTERN_ESCAPE": 14, - "CLI_DBMS_NAME": 17, - "CLI_DBMS_VER": 18, - "CLI_ACCESSIBLE_TABLES": 19, - "CLI_ACCESSIBLE_PROCEDURES": 20, - "CLI_CURSOR_COMMIT_BEHAVIOR": 23, - "CLI_DATA_SOURCE_READ_ONLY": 25, - "CLI_DEFAULT_TXN_ISOLATION": 26, - "CLI_IDENTIFIER_CASE": 28, - "CLI_IDENTIFIER_QUOTE_CHAR": 29, - "CLI_MAX_COLUMN_NAME_LEN": 30, - "CLI_MAX_CURSOR_NAME_LEN": 31, - "CLI_MAX_SCHEMA_NAME_LEN": 32, - "CLI_MAX_CATALOG_NAME_LEN": 34, - "CLI_MAX_TABLE_NAME_LEN": 35, - "CLI_SCROLL_CONCURRENCY": 43, - "CLI_TXN_CAPABLE": 46, - "CLI_USER_NAME": 47, - "CLI_TXN_ISOLATION_OPTION": 72, - "CLI_INTEGRITY": 73, - "CLI_GETDATA_EXTENSIONS": 81, - "CLI_NULL_COLLATION": 85, - "CLI_ALTER_TABLE": 86, - "CLI_ORDER_BY_COLUMNS_IN_SELECT": 90, - "CLI_SPECIAL_CHARACTERS": 94, - "CLI_MAX_COLUMNS_IN_GROUP_BY": 97, - "CLI_MAX_COLUMNS_IN_INDEX": 98, - "CLI_MAX_COLUMNS_IN_ORDER_BY": 99, - "CLI_MAX_COLUMNS_IN_SELECT": 100, - "CLI_MAX_COLUMNS_IN_TABLE": 101, - "CLI_MAX_INDEX_SIZE": 102, - "CLI_MAX_ROW_SIZE": 104, - "CLI_MAX_STATEMENT_LEN": 105, - "CLI_MAX_TABLES_IN_SELECT": 106, - "CLI_MAX_USER_NAME_LEN": 107, - "CLI_OJ_CAPABILITIES": 115, - "CLI_XOPEN_CLI_YEAR": 10000, - "CLI_CURSOR_SENSITIVITY": 10001, - "CLI_DESCRIBE_PARAMETER": 10002, - "CLI_CATALOG_NAME": 10003, - "CLI_COLLATION_SEQ": 10004, - "CLI_MAX_IDENTIFIER_LEN": 10005, - } - - -class TFetchOrientation(object): - FETCH_NEXT = 0 - FETCH_PRIOR = 1 - FETCH_RELATIVE = 2 - FETCH_ABSOLUTE = 3 - FETCH_FIRST = 4 - FETCH_LAST = 5 - - _VALUES_TO_NAMES = { - 0: "FETCH_NEXT", - 1: "FETCH_PRIOR", - 2: "FETCH_RELATIVE", - 3: "FETCH_ABSOLUTE", - 4: "FETCH_FIRST", - 5: "FETCH_LAST", - } - - _NAMES_TO_VALUES = { - "FETCH_NEXT": 0, - "FETCH_PRIOR": 1, - "FETCH_RELATIVE": 2, - "FETCH_ABSOLUTE": 3, - "FETCH_FIRST": 4, - "FETCH_LAST": 5, - } - - -class TJobExecutionStatus(object): - IN_PROGRESS = 0 - COMPLETE = 1 - NOT_AVAILABLE = 2 - - _VALUES_TO_NAMES = { - 0: "IN_PROGRESS", - 1: "COMPLETE", - 2: "NOT_AVAILABLE", - } - - _NAMES_TO_VALUES = { - "IN_PROGRESS": 0, - "COMPLETE": 1, - "NOT_AVAILABLE": 2, - } - - -class TTypeQualifierValue(object): - """ - Attributes: - - i32Value - - stringValue - """ - - thrift_spec = ( - None, # 0 - (1, TType.I32, 'i32Value', None, None, ), # 1 - (2, TType.STRING, 'stringValue', 'UTF8', None, ), # 2 - ) - - def __init__(self, i32Value=None, stringValue=None,): - self.i32Value = i32Value - self.stringValue = stringValue - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.I32: - self.i32Value = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.stringValue = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TTypeQualifierValue') - if self.i32Value is not None: - oprot.writeFieldBegin('i32Value', TType.I32, 1) - oprot.writeI32(self.i32Value) - oprot.writeFieldEnd() - if self.stringValue is not None: - oprot.writeFieldBegin('stringValue', TType.STRING, 2) - oprot.writeString(self.stringValue.encode('utf-8') if sys.version_info[0] == 2 else self.stringValue) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TTypeQualifiers(object): - """ - Attributes: - - qualifiers - """ - - thrift_spec = ( - None, # 0 - (1, TType.MAP, 'qualifiers', (TType.STRING, 'UTF8', TType.STRUCT, (TTypeQualifierValue, TTypeQualifierValue.thrift_spec), False), None, ), # 1 - ) - - def __init__(self, qualifiers=None,): - self.qualifiers = qualifiers - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.MAP: - self.qualifiers = {} - (_ktype1, _vtype2, _size0) = iprot.readMapBegin() - for _i4 in range(_size0): - _key5 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - _val6 = TTypeQualifierValue() - _val6.read(iprot) - self.qualifiers[_key5] = _val6 - iprot.readMapEnd() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TTypeQualifiers') - if self.qualifiers is not None: - oprot.writeFieldBegin('qualifiers', TType.MAP, 1) - oprot.writeMapBegin(TType.STRING, TType.STRUCT, len(self.qualifiers)) - for kiter7, viter8 in self.qualifiers.items(): - oprot.writeString(kiter7.encode('utf-8') if sys.version_info[0] == 2 else kiter7) - viter8.write(oprot) - oprot.writeMapEnd() - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.qualifiers is None: - raise TProtocolException(message='Required field qualifiers is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TPrimitiveTypeEntry(object): - """ - Attributes: - - type - - typeQualifiers - """ - - thrift_spec = ( - None, # 0 - (1, TType.I32, 'type', None, None, ), # 1 - (2, TType.STRUCT, 'typeQualifiers', (TTypeQualifiers, TTypeQualifiers.thrift_spec), None, ), # 2 - ) - - def __init__(self, type=None, typeQualifiers=None,): - self.type = type - self.typeQualifiers = typeQualifiers - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.I32: - self.type = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.typeQualifiers = TTypeQualifiers() - self.typeQualifiers.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TPrimitiveTypeEntry') - if self.type is not None: - oprot.writeFieldBegin('type', TType.I32, 1) - oprot.writeI32(self.type) - oprot.writeFieldEnd() - if self.typeQualifiers is not None: - oprot.writeFieldBegin('typeQualifiers', TType.STRUCT, 2) - self.typeQualifiers.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.type is None: - raise TProtocolException(message='Required field type is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TArrayTypeEntry(object): - """ - Attributes: - - objectTypePtr - """ - - thrift_spec = ( - None, # 0 - (1, TType.I32, 'objectTypePtr', None, None, ), # 1 - ) - - def __init__(self, objectTypePtr=None,): - self.objectTypePtr = objectTypePtr - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.I32: - self.objectTypePtr = iprot.readI32() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TArrayTypeEntry') - if self.objectTypePtr is not None: - oprot.writeFieldBegin('objectTypePtr', TType.I32, 1) - oprot.writeI32(self.objectTypePtr) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.objectTypePtr is None: - raise TProtocolException(message='Required field objectTypePtr is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TMapTypeEntry(object): - """ - Attributes: - - keyTypePtr - - valueTypePtr - """ - - thrift_spec = ( - None, # 0 - (1, TType.I32, 'keyTypePtr', None, None, ), # 1 - (2, TType.I32, 'valueTypePtr', None, None, ), # 2 - ) - - def __init__(self, keyTypePtr=None, valueTypePtr=None,): - self.keyTypePtr = keyTypePtr - self.valueTypePtr = valueTypePtr - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.I32: - self.keyTypePtr = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.I32: - self.valueTypePtr = iprot.readI32() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TMapTypeEntry') - if self.keyTypePtr is not None: - oprot.writeFieldBegin('keyTypePtr', TType.I32, 1) - oprot.writeI32(self.keyTypePtr) - oprot.writeFieldEnd() - if self.valueTypePtr is not None: - oprot.writeFieldBegin('valueTypePtr', TType.I32, 2) - oprot.writeI32(self.valueTypePtr) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.keyTypePtr is None: - raise TProtocolException(message='Required field keyTypePtr is unset!') - if self.valueTypePtr is None: - raise TProtocolException(message='Required field valueTypePtr is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TStructTypeEntry(object): - """ - Attributes: - - nameToTypePtr - """ - - thrift_spec = ( - None, # 0 - (1, TType.MAP, 'nameToTypePtr', (TType.STRING, 'UTF8', TType.I32, None, False), None, ), # 1 - ) - - def __init__(self, nameToTypePtr=None,): - self.nameToTypePtr = nameToTypePtr - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.MAP: - self.nameToTypePtr = {} - (_ktype10, _vtype11, _size9) = iprot.readMapBegin() - for _i13 in range(_size9): - _key14 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - _val15 = iprot.readI32() - self.nameToTypePtr[_key14] = _val15 - iprot.readMapEnd() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TStructTypeEntry') - if self.nameToTypePtr is not None: - oprot.writeFieldBegin('nameToTypePtr', TType.MAP, 1) - oprot.writeMapBegin(TType.STRING, TType.I32, len(self.nameToTypePtr)) - for kiter16, viter17 in self.nameToTypePtr.items(): - oprot.writeString(kiter16.encode('utf-8') if sys.version_info[0] == 2 else kiter16) - oprot.writeI32(viter17) - oprot.writeMapEnd() - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.nameToTypePtr is None: - raise TProtocolException(message='Required field nameToTypePtr is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TUnionTypeEntry(object): - """ - Attributes: - - nameToTypePtr - """ - - thrift_spec = ( - None, # 0 - (1, TType.MAP, 'nameToTypePtr', (TType.STRING, 'UTF8', TType.I32, None, False), None, ), # 1 - ) - - def __init__(self, nameToTypePtr=None,): - self.nameToTypePtr = nameToTypePtr - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.MAP: - self.nameToTypePtr = {} - (_ktype19, _vtype20, _size18) = iprot.readMapBegin() - for _i22 in range(_size18): - _key23 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - _val24 = iprot.readI32() - self.nameToTypePtr[_key23] = _val24 - iprot.readMapEnd() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TUnionTypeEntry') - if self.nameToTypePtr is not None: - oprot.writeFieldBegin('nameToTypePtr', TType.MAP, 1) - oprot.writeMapBegin(TType.STRING, TType.I32, len(self.nameToTypePtr)) - for kiter25, viter26 in self.nameToTypePtr.items(): - oprot.writeString(kiter25.encode('utf-8') if sys.version_info[0] == 2 else kiter25) - oprot.writeI32(viter26) - oprot.writeMapEnd() - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.nameToTypePtr is None: - raise TProtocolException(message='Required field nameToTypePtr is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TUserDefinedTypeEntry(object): - """ - Attributes: - - typeClassName - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRING, 'typeClassName', 'UTF8', None, ), # 1 - ) - - def __init__(self, typeClassName=None,): - self.typeClassName = typeClassName - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRING: - self.typeClassName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TUserDefinedTypeEntry') - if self.typeClassName is not None: - oprot.writeFieldBegin('typeClassName', TType.STRING, 1) - oprot.writeString(self.typeClassName.encode('utf-8') if sys.version_info[0] == 2 else self.typeClassName) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.typeClassName is None: - raise TProtocolException(message='Required field typeClassName is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TTypeEntry(object): - """ - Attributes: - - primitiveEntry - - arrayEntry - - mapEntry - - structEntry - - unionEntry - - userDefinedTypeEntry - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'primitiveEntry', (TPrimitiveTypeEntry, TPrimitiveTypeEntry.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'arrayEntry', (TArrayTypeEntry, TArrayTypeEntry.thrift_spec), None, ), # 2 - (3, TType.STRUCT, 'mapEntry', (TMapTypeEntry, TMapTypeEntry.thrift_spec), None, ), # 3 - (4, TType.STRUCT, 'structEntry', (TStructTypeEntry, TStructTypeEntry.thrift_spec), None, ), # 4 - (5, TType.STRUCT, 'unionEntry', (TUnionTypeEntry, TUnionTypeEntry.thrift_spec), None, ), # 5 - (6, TType.STRUCT, 'userDefinedTypeEntry', (TUserDefinedTypeEntry, TUserDefinedTypeEntry.thrift_spec), None, ), # 6 - ) - - def __init__(self, primitiveEntry=None, arrayEntry=None, mapEntry=None, structEntry=None, unionEntry=None, userDefinedTypeEntry=None,): - self.primitiveEntry = primitiveEntry - self.arrayEntry = arrayEntry - self.mapEntry = mapEntry - self.structEntry = structEntry - self.unionEntry = unionEntry - self.userDefinedTypeEntry = userDefinedTypeEntry - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.primitiveEntry = TPrimitiveTypeEntry() - self.primitiveEntry.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.arrayEntry = TArrayTypeEntry() - self.arrayEntry.read(iprot) - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRUCT: - self.mapEntry = TMapTypeEntry() - self.mapEntry.read(iprot) - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.STRUCT: - self.structEntry = TStructTypeEntry() - self.structEntry.read(iprot) - else: - iprot.skip(ftype) - elif fid == 5: - if ftype == TType.STRUCT: - self.unionEntry = TUnionTypeEntry() - self.unionEntry.read(iprot) - else: - iprot.skip(ftype) - elif fid == 6: - if ftype == TType.STRUCT: - self.userDefinedTypeEntry = TUserDefinedTypeEntry() - self.userDefinedTypeEntry.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TTypeEntry') - if self.primitiveEntry is not None: - oprot.writeFieldBegin('primitiveEntry', TType.STRUCT, 1) - self.primitiveEntry.write(oprot) - oprot.writeFieldEnd() - if self.arrayEntry is not None: - oprot.writeFieldBegin('arrayEntry', TType.STRUCT, 2) - self.arrayEntry.write(oprot) - oprot.writeFieldEnd() - if self.mapEntry is not None: - oprot.writeFieldBegin('mapEntry', TType.STRUCT, 3) - self.mapEntry.write(oprot) - oprot.writeFieldEnd() - if self.structEntry is not None: - oprot.writeFieldBegin('structEntry', TType.STRUCT, 4) - self.structEntry.write(oprot) - oprot.writeFieldEnd() - if self.unionEntry is not None: - oprot.writeFieldBegin('unionEntry', TType.STRUCT, 5) - self.unionEntry.write(oprot) - oprot.writeFieldEnd() - if self.userDefinedTypeEntry is not None: - oprot.writeFieldBegin('userDefinedTypeEntry', TType.STRUCT, 6) - self.userDefinedTypeEntry.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TTypeDesc(object): - """ - Attributes: - - types - """ - - thrift_spec = ( - None, # 0 - (1, TType.LIST, 'types', (TType.STRUCT, (TTypeEntry, TTypeEntry.thrift_spec), False), None, ), # 1 - ) - - def __init__(self, types=None,): - self.types = types - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.LIST: - self.types = [] - (_etype30, _size27) = iprot.readListBegin() - for _i31 in range(_size27): - _elem32 = TTypeEntry() - _elem32.read(iprot) - self.types.append(_elem32) - iprot.readListEnd() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TTypeDesc') - if self.types is not None: - oprot.writeFieldBegin('types', TType.LIST, 1) - oprot.writeListBegin(TType.STRUCT, len(self.types)) - for iter33 in self.types: - iter33.write(oprot) - oprot.writeListEnd() - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.types is None: - raise TProtocolException(message='Required field types is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TColumnDesc(object): - """ - Attributes: - - columnName - - typeDesc - - position - - comment - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRING, 'columnName', 'UTF8', None, ), # 1 - (2, TType.STRUCT, 'typeDesc', (TTypeDesc, TTypeDesc.thrift_spec), None, ), # 2 - (3, TType.I32, 'position', None, None, ), # 3 - (4, TType.STRING, 'comment', 'UTF8', None, ), # 4 - ) - - def __init__(self, columnName=None, typeDesc=None, position=None, comment=None,): - self.columnName = columnName - self.typeDesc = typeDesc - self.position = position - self.comment = comment - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRING: - self.columnName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.typeDesc = TTypeDesc() - self.typeDesc.read(iprot) - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.I32: - self.position = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.STRING: - self.comment = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TColumnDesc') - if self.columnName is not None: - oprot.writeFieldBegin('columnName', TType.STRING, 1) - oprot.writeString(self.columnName.encode('utf-8') if sys.version_info[0] == 2 else self.columnName) - oprot.writeFieldEnd() - if self.typeDesc is not None: - oprot.writeFieldBegin('typeDesc', TType.STRUCT, 2) - self.typeDesc.write(oprot) - oprot.writeFieldEnd() - if self.position is not None: - oprot.writeFieldBegin('position', TType.I32, 3) - oprot.writeI32(self.position) - oprot.writeFieldEnd() - if self.comment is not None: - oprot.writeFieldBegin('comment', TType.STRING, 4) - oprot.writeString(self.comment.encode('utf-8') if sys.version_info[0] == 2 else self.comment) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.columnName is None: - raise TProtocolException(message='Required field columnName is unset!') - if self.typeDesc is None: - raise TProtocolException(message='Required field typeDesc is unset!') - if self.position is None: - raise TProtocolException(message='Required field position is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TTableSchema(object): - """ - Attributes: - - columns - """ - - thrift_spec = ( - None, # 0 - (1, TType.LIST, 'columns', (TType.STRUCT, (TColumnDesc, TColumnDesc.thrift_spec), False), None, ), # 1 - ) - - def __init__(self, columns=None,): - self.columns = columns - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.LIST: - self.columns = [] - (_etype37, _size34) = iprot.readListBegin() - for _i38 in range(_size34): - _elem39 = TColumnDesc() - _elem39.read(iprot) - self.columns.append(_elem39) - iprot.readListEnd() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TTableSchema') - if self.columns is not None: - oprot.writeFieldBegin('columns', TType.LIST, 1) - oprot.writeListBegin(TType.STRUCT, len(self.columns)) - for iter40 in self.columns: - iter40.write(oprot) - oprot.writeListEnd() - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.columns is None: - raise TProtocolException(message='Required field columns is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TBoolValue(object): - """ - Attributes: - - value - """ - - thrift_spec = ( - None, # 0 - (1, TType.BOOL, 'value', None, None, ), # 1 - ) - - def __init__(self, value=None,): - self.value = value - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.BOOL: - self.value = iprot.readBool() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TBoolValue') - if self.value is not None: - oprot.writeFieldBegin('value', TType.BOOL, 1) - oprot.writeBool(self.value) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TByteValue(object): - """ - Attributes: - - value - """ - - thrift_spec = ( - None, # 0 - (1, TType.BYTE, 'value', None, None, ), # 1 - ) - - def __init__(self, value=None,): - self.value = value - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.BYTE: - self.value = iprot.readByte() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TByteValue') - if self.value is not None: - oprot.writeFieldBegin('value', TType.BYTE, 1) - oprot.writeByte(self.value) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TI16Value(object): - """ - Attributes: - - value - """ - - thrift_spec = ( - None, # 0 - (1, TType.I16, 'value', None, None, ), # 1 - ) - - def __init__(self, value=None,): - self.value = value - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.I16: - self.value = iprot.readI16() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TI16Value') - if self.value is not None: - oprot.writeFieldBegin('value', TType.I16, 1) - oprot.writeI16(self.value) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TI32Value(object): - """ - Attributes: - - value - """ - - thrift_spec = ( - None, # 0 - (1, TType.I32, 'value', None, None, ), # 1 - ) - - def __init__(self, value=None,): - self.value = value - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.I32: - self.value = iprot.readI32() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TI32Value') - if self.value is not None: - oprot.writeFieldBegin('value', TType.I32, 1) - oprot.writeI32(self.value) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TI64Value(object): - """ - Attributes: - - value - """ - - thrift_spec = ( - None, # 0 - (1, TType.I64, 'value', None, None, ), # 1 - ) - - def __init__(self, value=None,): - self.value = value - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.I64: - self.value = iprot.readI64() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TI64Value') - if self.value is not None: - oprot.writeFieldBegin('value', TType.I64, 1) - oprot.writeI64(self.value) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TDoubleValue(object): - """ - Attributes: - - value - """ - - thrift_spec = ( - None, # 0 - (1, TType.DOUBLE, 'value', None, None, ), # 1 - ) - - def __init__(self, value=None,): - self.value = value - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.DOUBLE: - self.value = iprot.readDouble() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TDoubleValue') - if self.value is not None: - oprot.writeFieldBegin('value', TType.DOUBLE, 1) - oprot.writeDouble(self.value) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TStringValue(object): - """ - Attributes: - - value - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRING, 'value', 'UTF8', None, ), # 1 - ) - - def __init__(self, value=None,): - self.value = value - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRING: - self.value = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TStringValue') - if self.value is not None: - oprot.writeFieldBegin('value', TType.STRING, 1) - oprot.writeString(self.value.encode('utf-8') if sys.version_info[0] == 2 else self.value) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TColumnValue(object): - """ - Attributes: - - boolVal - - byteVal - - i16Val - - i32Val - - i64Val - - doubleVal - - stringVal - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'boolVal', (TBoolValue, TBoolValue.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'byteVal', (TByteValue, TByteValue.thrift_spec), None, ), # 2 - (3, TType.STRUCT, 'i16Val', (TI16Value, TI16Value.thrift_spec), None, ), # 3 - (4, TType.STRUCT, 'i32Val', (TI32Value, TI32Value.thrift_spec), None, ), # 4 - (5, TType.STRUCT, 'i64Val', (TI64Value, TI64Value.thrift_spec), None, ), # 5 - (6, TType.STRUCT, 'doubleVal', (TDoubleValue, TDoubleValue.thrift_spec), None, ), # 6 - (7, TType.STRUCT, 'stringVal', (TStringValue, TStringValue.thrift_spec), None, ), # 7 - ) - - def __init__(self, boolVal=None, byteVal=None, i16Val=None, i32Val=None, i64Val=None, doubleVal=None, stringVal=None,): - self.boolVal = boolVal - self.byteVal = byteVal - self.i16Val = i16Val - self.i32Val = i32Val - self.i64Val = i64Val - self.doubleVal = doubleVal - self.stringVal = stringVal - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.boolVal = TBoolValue() - self.boolVal.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.byteVal = TByteValue() - self.byteVal.read(iprot) - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRUCT: - self.i16Val = TI16Value() - self.i16Val.read(iprot) - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.STRUCT: - self.i32Val = TI32Value() - self.i32Val.read(iprot) - else: - iprot.skip(ftype) - elif fid == 5: - if ftype == TType.STRUCT: - self.i64Val = TI64Value() - self.i64Val.read(iprot) - else: - iprot.skip(ftype) - elif fid == 6: - if ftype == TType.STRUCT: - self.doubleVal = TDoubleValue() - self.doubleVal.read(iprot) - else: - iprot.skip(ftype) - elif fid == 7: - if ftype == TType.STRUCT: - self.stringVal = TStringValue() - self.stringVal.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TColumnValue') - if self.boolVal is not None: - oprot.writeFieldBegin('boolVal', TType.STRUCT, 1) - self.boolVal.write(oprot) - oprot.writeFieldEnd() - if self.byteVal is not None: - oprot.writeFieldBegin('byteVal', TType.STRUCT, 2) - self.byteVal.write(oprot) - oprot.writeFieldEnd() - if self.i16Val is not None: - oprot.writeFieldBegin('i16Val', TType.STRUCT, 3) - self.i16Val.write(oprot) - oprot.writeFieldEnd() - if self.i32Val is not None: - oprot.writeFieldBegin('i32Val', TType.STRUCT, 4) - self.i32Val.write(oprot) - oprot.writeFieldEnd() - if self.i64Val is not None: - oprot.writeFieldBegin('i64Val', TType.STRUCT, 5) - self.i64Val.write(oprot) - oprot.writeFieldEnd() - if self.doubleVal is not None: - oprot.writeFieldBegin('doubleVal', TType.STRUCT, 6) - self.doubleVal.write(oprot) - oprot.writeFieldEnd() - if self.stringVal is not None: - oprot.writeFieldBegin('stringVal', TType.STRUCT, 7) - self.stringVal.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TRow(object): - """ - Attributes: - - colVals - """ - - thrift_spec = ( - None, # 0 - (1, TType.LIST, 'colVals', (TType.STRUCT, (TColumnValue, TColumnValue.thrift_spec), False), None, ), # 1 - ) - - def __init__(self, colVals=None,): - self.colVals = colVals - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.LIST: - self.colVals = [] - (_etype44, _size41) = iprot.readListBegin() - for _i45 in range(_size41): - _elem46 = TColumnValue() - _elem46.read(iprot) - self.colVals.append(_elem46) - iprot.readListEnd() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TRow') - if self.colVals is not None: - oprot.writeFieldBegin('colVals', TType.LIST, 1) - oprot.writeListBegin(TType.STRUCT, len(self.colVals)) - for iter47 in self.colVals: - iter47.write(oprot) - oprot.writeListEnd() - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.colVals is None: - raise TProtocolException(message='Required field colVals is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TBoolColumn(object): - """ - Attributes: - - values - - nulls - """ - - thrift_spec = ( - None, # 0 - (1, TType.LIST, 'values', (TType.BOOL, None, False), None, ), # 1 - (2, TType.STRING, 'nulls', 'BINARY', None, ), # 2 - ) - - def __init__(self, values=None, nulls=None,): - self.values = values - self.nulls = nulls - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.LIST: - self.values = [] - (_etype51, _size48) = iprot.readListBegin() - for _i52 in range(_size48): - _elem53 = iprot.readBool() - self.values.append(_elem53) - iprot.readListEnd() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.nulls = iprot.readBinary() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TBoolColumn') - if self.values is not None: - oprot.writeFieldBegin('values', TType.LIST, 1) - oprot.writeListBegin(TType.BOOL, len(self.values)) - for iter54 in self.values: - oprot.writeBool(iter54) - oprot.writeListEnd() - oprot.writeFieldEnd() - if self.nulls is not None: - oprot.writeFieldBegin('nulls', TType.STRING, 2) - oprot.writeBinary(self.nulls) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.values is None: - raise TProtocolException(message='Required field values is unset!') - if self.nulls is None: - raise TProtocolException(message='Required field nulls is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TByteColumn(object): - """ - Attributes: - - values - - nulls - """ - - thrift_spec = ( - None, # 0 - (1, TType.LIST, 'values', (TType.BYTE, None, False), None, ), # 1 - (2, TType.STRING, 'nulls', 'BINARY', None, ), # 2 - ) - - def __init__(self, values=None, nulls=None,): - self.values = values - self.nulls = nulls - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.LIST: - self.values = [] - (_etype58, _size55) = iprot.readListBegin() - for _i59 in range(_size55): - _elem60 = iprot.readByte() - self.values.append(_elem60) - iprot.readListEnd() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.nulls = iprot.readBinary() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TByteColumn') - if self.values is not None: - oprot.writeFieldBegin('values', TType.LIST, 1) - oprot.writeListBegin(TType.BYTE, len(self.values)) - for iter61 in self.values: - oprot.writeByte(iter61) - oprot.writeListEnd() - oprot.writeFieldEnd() - if self.nulls is not None: - oprot.writeFieldBegin('nulls', TType.STRING, 2) - oprot.writeBinary(self.nulls) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.values is None: - raise TProtocolException(message='Required field values is unset!') - if self.nulls is None: - raise TProtocolException(message='Required field nulls is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TI16Column(object): - """ - Attributes: - - values - - nulls - """ - - thrift_spec = ( - None, # 0 - (1, TType.LIST, 'values', (TType.I16, None, False), None, ), # 1 - (2, TType.STRING, 'nulls', 'BINARY', None, ), # 2 - ) - - def __init__(self, values=None, nulls=None,): - self.values = values - self.nulls = nulls - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.LIST: - self.values = [] - (_etype65, _size62) = iprot.readListBegin() - for _i66 in range(_size62): - _elem67 = iprot.readI16() - self.values.append(_elem67) - iprot.readListEnd() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.nulls = iprot.readBinary() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TI16Column') - if self.values is not None: - oprot.writeFieldBegin('values', TType.LIST, 1) - oprot.writeListBegin(TType.I16, len(self.values)) - for iter68 in self.values: - oprot.writeI16(iter68) - oprot.writeListEnd() - oprot.writeFieldEnd() - if self.nulls is not None: - oprot.writeFieldBegin('nulls', TType.STRING, 2) - oprot.writeBinary(self.nulls) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.values is None: - raise TProtocolException(message='Required field values is unset!') - if self.nulls is None: - raise TProtocolException(message='Required field nulls is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TI32Column(object): - """ - Attributes: - - values - - nulls - """ - - thrift_spec = ( - None, # 0 - (1, TType.LIST, 'values', (TType.I32, None, False), None, ), # 1 - (2, TType.STRING, 'nulls', 'BINARY', None, ), # 2 - ) - - def __init__(self, values=None, nulls=None,): - self.values = values - self.nulls = nulls - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.LIST: - self.values = [] - (_etype72, _size69) = iprot.readListBegin() - for _i73 in range(_size69): - _elem74 = iprot.readI32() - self.values.append(_elem74) - iprot.readListEnd() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.nulls = iprot.readBinary() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TI32Column') - if self.values is not None: - oprot.writeFieldBegin('values', TType.LIST, 1) - oprot.writeListBegin(TType.I32, len(self.values)) - for iter75 in self.values: - oprot.writeI32(iter75) - oprot.writeListEnd() - oprot.writeFieldEnd() - if self.nulls is not None: - oprot.writeFieldBegin('nulls', TType.STRING, 2) - oprot.writeBinary(self.nulls) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.values is None: - raise TProtocolException(message='Required field values is unset!') - if self.nulls is None: - raise TProtocolException(message='Required field nulls is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TI64Column(object): - """ - Attributes: - - values - - nulls - """ - - thrift_spec = ( - None, # 0 - (1, TType.LIST, 'values', (TType.I64, None, False), None, ), # 1 - (2, TType.STRING, 'nulls', 'BINARY', None, ), # 2 - ) - - def __init__(self, values=None, nulls=None,): - self.values = values - self.nulls = nulls - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.LIST: - self.values = [] - (_etype79, _size76) = iprot.readListBegin() - for _i80 in range(_size76): - _elem81 = iprot.readI64() - self.values.append(_elem81) - iprot.readListEnd() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.nulls = iprot.readBinary() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TI64Column') - if self.values is not None: - oprot.writeFieldBegin('values', TType.LIST, 1) - oprot.writeListBegin(TType.I64, len(self.values)) - for iter82 in self.values: - oprot.writeI64(iter82) - oprot.writeListEnd() - oprot.writeFieldEnd() - if self.nulls is not None: - oprot.writeFieldBegin('nulls', TType.STRING, 2) - oprot.writeBinary(self.nulls) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.values is None: - raise TProtocolException(message='Required field values is unset!') - if self.nulls is None: - raise TProtocolException(message='Required field nulls is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TDoubleColumn(object): - """ - Attributes: - - values - - nulls - """ - - thrift_spec = ( - None, # 0 - (1, TType.LIST, 'values', (TType.DOUBLE, None, False), None, ), # 1 - (2, TType.STRING, 'nulls', 'BINARY', None, ), # 2 - ) - - def __init__(self, values=None, nulls=None,): - self.values = values - self.nulls = nulls - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.LIST: - self.values = [] - (_etype86, _size83) = iprot.readListBegin() - for _i87 in range(_size83): - _elem88 = iprot.readDouble() - self.values.append(_elem88) - iprot.readListEnd() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.nulls = iprot.readBinary() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TDoubleColumn') - if self.values is not None: - oprot.writeFieldBegin('values', TType.LIST, 1) - oprot.writeListBegin(TType.DOUBLE, len(self.values)) - for iter89 in self.values: - oprot.writeDouble(iter89) - oprot.writeListEnd() - oprot.writeFieldEnd() - if self.nulls is not None: - oprot.writeFieldBegin('nulls', TType.STRING, 2) - oprot.writeBinary(self.nulls) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.values is None: - raise TProtocolException(message='Required field values is unset!') - if self.nulls is None: - raise TProtocolException(message='Required field nulls is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TStringColumn(object): - """ - Attributes: - - values - - nulls - """ - - thrift_spec = ( - None, # 0 - (1, TType.LIST, 'values', (TType.STRING, 'UTF8', False), None, ), # 1 - (2, TType.STRING, 'nulls', 'BINARY', None, ), # 2 - ) - - def __init__(self, values=None, nulls=None,): - self.values = values - self.nulls = nulls - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.LIST: - self.values = [] - (_etype93, _size90) = iprot.readListBegin() - for _i94 in range(_size90): - _elem95 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - self.values.append(_elem95) - iprot.readListEnd() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.nulls = iprot.readBinary() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TStringColumn') - if self.values is not None: - oprot.writeFieldBegin('values', TType.LIST, 1) - oprot.writeListBegin(TType.STRING, len(self.values)) - for iter96 in self.values: - oprot.writeString(iter96.encode('utf-8') if sys.version_info[0] == 2 else iter96) - oprot.writeListEnd() - oprot.writeFieldEnd() - if self.nulls is not None: - oprot.writeFieldBegin('nulls', TType.STRING, 2) - oprot.writeBinary(self.nulls) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.values is None: - raise TProtocolException(message='Required field values is unset!') - if self.nulls is None: - raise TProtocolException(message='Required field nulls is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TBinaryColumn(object): - """ - Attributes: - - values - - nulls - """ - - thrift_spec = ( - None, # 0 - (1, TType.LIST, 'values', (TType.STRING, 'BINARY', False), None, ), # 1 - (2, TType.STRING, 'nulls', 'BINARY', None, ), # 2 - ) - - def __init__(self, values=None, nulls=None,): - self.values = values - self.nulls = nulls - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.LIST: - self.values = [] - (_etype100, _size97) = iprot.readListBegin() - for _i101 in range(_size97): - _elem102 = iprot.readBinary() - self.values.append(_elem102) - iprot.readListEnd() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.nulls = iprot.readBinary() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TBinaryColumn') - if self.values is not None: - oprot.writeFieldBegin('values', TType.LIST, 1) - oprot.writeListBegin(TType.STRING, len(self.values)) - for iter103 in self.values: - oprot.writeBinary(iter103) - oprot.writeListEnd() - oprot.writeFieldEnd() - if self.nulls is not None: - oprot.writeFieldBegin('nulls', TType.STRING, 2) - oprot.writeBinary(self.nulls) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.values is None: - raise TProtocolException(message='Required field values is unset!') - if self.nulls is None: - raise TProtocolException(message='Required field nulls is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TColumn(object): - """ - Attributes: - - boolVal - - byteVal - - i16Val - - i32Val - - i64Val - - doubleVal - - stringVal - - binaryVal - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'boolVal', (TBoolColumn, TBoolColumn.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'byteVal', (TByteColumn, TByteColumn.thrift_spec), None, ), # 2 - (3, TType.STRUCT, 'i16Val', (TI16Column, TI16Column.thrift_spec), None, ), # 3 - (4, TType.STRUCT, 'i32Val', (TI32Column, TI32Column.thrift_spec), None, ), # 4 - (5, TType.STRUCT, 'i64Val', (TI64Column, TI64Column.thrift_spec), None, ), # 5 - (6, TType.STRUCT, 'doubleVal', (TDoubleColumn, TDoubleColumn.thrift_spec), None, ), # 6 - (7, TType.STRUCT, 'stringVal', (TStringColumn, TStringColumn.thrift_spec), None, ), # 7 - (8, TType.STRUCT, 'binaryVal', (TBinaryColumn, TBinaryColumn.thrift_spec), None, ), # 8 - ) - - def __init__(self, boolVal=None, byteVal=None, i16Val=None, i32Val=None, i64Val=None, doubleVal=None, stringVal=None, binaryVal=None,): - self.boolVal = boolVal - self.byteVal = byteVal - self.i16Val = i16Val - self.i32Val = i32Val - self.i64Val = i64Val - self.doubleVal = doubleVal - self.stringVal = stringVal - self.binaryVal = binaryVal - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.boolVal = TBoolColumn() - self.boolVal.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.byteVal = TByteColumn() - self.byteVal.read(iprot) - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRUCT: - self.i16Val = TI16Column() - self.i16Val.read(iprot) - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.STRUCT: - self.i32Val = TI32Column() - self.i32Val.read(iprot) - else: - iprot.skip(ftype) - elif fid == 5: - if ftype == TType.STRUCT: - self.i64Val = TI64Column() - self.i64Val.read(iprot) - else: - iprot.skip(ftype) - elif fid == 6: - if ftype == TType.STRUCT: - self.doubleVal = TDoubleColumn() - self.doubleVal.read(iprot) - else: - iprot.skip(ftype) - elif fid == 7: - if ftype == TType.STRUCT: - self.stringVal = TStringColumn() - self.stringVal.read(iprot) - else: - iprot.skip(ftype) - elif fid == 8: - if ftype == TType.STRUCT: - self.binaryVal = TBinaryColumn() - self.binaryVal.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TColumn') - if self.boolVal is not None: - oprot.writeFieldBegin('boolVal', TType.STRUCT, 1) - self.boolVal.write(oprot) - oprot.writeFieldEnd() - if self.byteVal is not None: - oprot.writeFieldBegin('byteVal', TType.STRUCT, 2) - self.byteVal.write(oprot) - oprot.writeFieldEnd() - if self.i16Val is not None: - oprot.writeFieldBegin('i16Val', TType.STRUCT, 3) - self.i16Val.write(oprot) - oprot.writeFieldEnd() - if self.i32Val is not None: - oprot.writeFieldBegin('i32Val', TType.STRUCT, 4) - self.i32Val.write(oprot) - oprot.writeFieldEnd() - if self.i64Val is not None: - oprot.writeFieldBegin('i64Val', TType.STRUCT, 5) - self.i64Val.write(oprot) - oprot.writeFieldEnd() - if self.doubleVal is not None: - oprot.writeFieldBegin('doubleVal', TType.STRUCT, 6) - self.doubleVal.write(oprot) - oprot.writeFieldEnd() - if self.stringVal is not None: - oprot.writeFieldBegin('stringVal', TType.STRUCT, 7) - self.stringVal.write(oprot) - oprot.writeFieldEnd() - if self.binaryVal is not None: - oprot.writeFieldBegin('binaryVal', TType.STRUCT, 8) - self.binaryVal.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TRowSet(object): - """ - Attributes: - - startRowOffset - - rows - - columns - - binaryColumns - - columnCount - """ - - thrift_spec = ( - None, # 0 - (1, TType.I64, 'startRowOffset', None, None, ), # 1 - (2, TType.LIST, 'rows', (TType.STRUCT, (TRow, TRow.thrift_spec), False), None, ), # 2 - (3, TType.LIST, 'columns', (TType.STRUCT, (TColumn, TColumn.thrift_spec), False), None, ), # 3 - (4, TType.STRING, 'binaryColumns', 'BINARY', None, ), # 4 - (5, TType.I32, 'columnCount', None, None, ), # 5 - ) - - def __init__(self, startRowOffset=None, rows=None, columns=None, binaryColumns=None, columnCount=None,): - self.startRowOffset = startRowOffset - self.rows = rows - self.columns = columns - self.binaryColumns = binaryColumns - self.columnCount = columnCount - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.I64: - self.startRowOffset = iprot.readI64() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.LIST: - self.rows = [] - (_etype107, _size104) = iprot.readListBegin() - for _i108 in range(_size104): - _elem109 = TRow() - _elem109.read(iprot) - self.rows.append(_elem109) - iprot.readListEnd() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.LIST: - self.columns = [] - (_etype113, _size110) = iprot.readListBegin() - for _i114 in range(_size110): - _elem115 = TColumn() - _elem115.read(iprot) - self.columns.append(_elem115) - iprot.readListEnd() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.STRING: - self.binaryColumns = iprot.readBinary() - else: - iprot.skip(ftype) - elif fid == 5: - if ftype == TType.I32: - self.columnCount = iprot.readI32() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TRowSet') - if self.startRowOffset is not None: - oprot.writeFieldBegin('startRowOffset', TType.I64, 1) - oprot.writeI64(self.startRowOffset) - oprot.writeFieldEnd() - if self.rows is not None: - oprot.writeFieldBegin('rows', TType.LIST, 2) - oprot.writeListBegin(TType.STRUCT, len(self.rows)) - for iter116 in self.rows: - iter116.write(oprot) - oprot.writeListEnd() - oprot.writeFieldEnd() - if self.columns is not None: - oprot.writeFieldBegin('columns', TType.LIST, 3) - oprot.writeListBegin(TType.STRUCT, len(self.columns)) - for iter117 in self.columns: - iter117.write(oprot) - oprot.writeListEnd() - oprot.writeFieldEnd() - if self.binaryColumns is not None: - oprot.writeFieldBegin('binaryColumns', TType.STRING, 4) - oprot.writeBinary(self.binaryColumns) - oprot.writeFieldEnd() - if self.columnCount is not None: - oprot.writeFieldBegin('columnCount', TType.I32, 5) - oprot.writeI32(self.columnCount) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.startRowOffset is None: - raise TProtocolException(message='Required field startRowOffset is unset!') - if self.rows is None: - raise TProtocolException(message='Required field rows is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TStatus(object): - """ - Attributes: - - statusCode - - infoMessages - - sqlState - - errorCode - - errorMessage - """ - - thrift_spec = ( - None, # 0 - (1, TType.I32, 'statusCode', None, None, ), # 1 - (2, TType.LIST, 'infoMessages', (TType.STRING, 'UTF8', False), None, ), # 2 - (3, TType.STRING, 'sqlState', 'UTF8', None, ), # 3 - (4, TType.I32, 'errorCode', None, None, ), # 4 - (5, TType.STRING, 'errorMessage', 'UTF8', None, ), # 5 - ) - - def __init__(self, statusCode=None, infoMessages=None, sqlState=None, errorCode=None, errorMessage=None,): - self.statusCode = statusCode - self.infoMessages = infoMessages - self.sqlState = sqlState - self.errorCode = errorCode - self.errorMessage = errorMessage - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.I32: - self.statusCode = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.LIST: - self.infoMessages = [] - (_etype121, _size118) = iprot.readListBegin() - for _i122 in range(_size118): - _elem123 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - self.infoMessages.append(_elem123) - iprot.readListEnd() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRING: - self.sqlState = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.I32: - self.errorCode = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 5: - if ftype == TType.STRING: - self.errorMessage = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TStatus') - if self.statusCode is not None: - oprot.writeFieldBegin('statusCode', TType.I32, 1) - oprot.writeI32(self.statusCode) - oprot.writeFieldEnd() - if self.infoMessages is not None: - oprot.writeFieldBegin('infoMessages', TType.LIST, 2) - oprot.writeListBegin(TType.STRING, len(self.infoMessages)) - for iter124 in self.infoMessages: - oprot.writeString(iter124.encode('utf-8') if sys.version_info[0] == 2 else iter124) - oprot.writeListEnd() - oprot.writeFieldEnd() - if self.sqlState is not None: - oprot.writeFieldBegin('sqlState', TType.STRING, 3) - oprot.writeString(self.sqlState.encode('utf-8') if sys.version_info[0] == 2 else self.sqlState) - oprot.writeFieldEnd() - if self.errorCode is not None: - oprot.writeFieldBegin('errorCode', TType.I32, 4) - oprot.writeI32(self.errorCode) - oprot.writeFieldEnd() - if self.errorMessage is not None: - oprot.writeFieldBegin('errorMessage', TType.STRING, 5) - oprot.writeString(self.errorMessage.encode('utf-8') if sys.version_info[0] == 2 else self.errorMessage) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.statusCode is None: - raise TProtocolException(message='Required field statusCode is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class THandleIdentifier(object): - """ - Attributes: - - guid - - secret - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRING, 'guid', 'BINARY', None, ), # 1 - (2, TType.STRING, 'secret', 'BINARY', None, ), # 2 - ) - - def __init__(self, guid=None, secret=None,): - self.guid = guid - self.secret = secret - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRING: - self.guid = iprot.readBinary() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.secret = iprot.readBinary() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('THandleIdentifier') - if self.guid is not None: - oprot.writeFieldBegin('guid', TType.STRING, 1) - oprot.writeBinary(self.guid) - oprot.writeFieldEnd() - if self.secret is not None: - oprot.writeFieldBegin('secret', TType.STRING, 2) - oprot.writeBinary(self.secret) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.guid is None: - raise TProtocolException(message='Required field guid is unset!') - if self.secret is None: - raise TProtocolException(message='Required field secret is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TSessionHandle(object): - """ - Attributes: - - sessionId - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionId', (THandleIdentifier, THandleIdentifier.thrift_spec), None, ), # 1 - ) - - def __init__(self, sessionId=None,): - self.sessionId = sessionId - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionId = THandleIdentifier() - self.sessionId.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TSessionHandle') - if self.sessionId is not None: - oprot.writeFieldBegin('sessionId', TType.STRUCT, 1) - self.sessionId.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionId is None: - raise TProtocolException(message='Required field sessionId is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TOperationHandle(object): - """ - Attributes: - - operationId - - operationType - - hasResultSet - - modifiedRowCount - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'operationId', (THandleIdentifier, THandleIdentifier.thrift_spec), None, ), # 1 - (2, TType.I32, 'operationType', None, None, ), # 2 - (3, TType.BOOL, 'hasResultSet', None, None, ), # 3 - (4, TType.DOUBLE, 'modifiedRowCount', None, None, ), # 4 - ) - - def __init__(self, operationId=None, operationType=None, hasResultSet=None, modifiedRowCount=None,): - self.operationId = operationId - self.operationType = operationType - self.hasResultSet = hasResultSet - self.modifiedRowCount = modifiedRowCount - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.operationId = THandleIdentifier() - self.operationId.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.I32: - self.operationType = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.BOOL: - self.hasResultSet = iprot.readBool() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.DOUBLE: - self.modifiedRowCount = iprot.readDouble() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TOperationHandle') - if self.operationId is not None: - oprot.writeFieldBegin('operationId', TType.STRUCT, 1) - self.operationId.write(oprot) - oprot.writeFieldEnd() - if self.operationType is not None: - oprot.writeFieldBegin('operationType', TType.I32, 2) - oprot.writeI32(self.operationType) - oprot.writeFieldEnd() - if self.hasResultSet is not None: - oprot.writeFieldBegin('hasResultSet', TType.BOOL, 3) - oprot.writeBool(self.hasResultSet) - oprot.writeFieldEnd() - if self.modifiedRowCount is not None: - oprot.writeFieldBegin('modifiedRowCount', TType.DOUBLE, 4) - oprot.writeDouble(self.modifiedRowCount) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.operationId is None: - raise TProtocolException(message='Required field operationId is unset!') - if self.operationType is None: - raise TProtocolException(message='Required field operationType is unset!') - if self.hasResultSet is None: - raise TProtocolException(message='Required field hasResultSet is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TOpenSessionReq(object): - """ - Attributes: - - client_protocol - - username - - password - - configuration - """ - - thrift_spec = ( - None, # 0 - (1, TType.I32, 'client_protocol', None, 9, ), # 1 - (2, TType.STRING, 'username', 'UTF8', None, ), # 2 - (3, TType.STRING, 'password', 'UTF8', None, ), # 3 - (4, TType.MAP, 'configuration', (TType.STRING, 'UTF8', TType.STRING, 'UTF8', False), None, ), # 4 - ) - - def __init__(self, client_protocol=thrift_spec[1][4], username=None, password=None, configuration=None,): - self.client_protocol = client_protocol - self.username = username - self.password = password - self.configuration = configuration - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.I32: - self.client_protocol = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.username = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRING: - self.password = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.MAP: - self.configuration = {} - (_ktype126, _vtype127, _size125) = iprot.readMapBegin() - for _i129 in range(_size125): - _key130 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - _val131 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - self.configuration[_key130] = _val131 - iprot.readMapEnd() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TOpenSessionReq') - if self.client_protocol is not None: - oprot.writeFieldBegin('client_protocol', TType.I32, 1) - oprot.writeI32(self.client_protocol) - oprot.writeFieldEnd() - if self.username is not None: - oprot.writeFieldBegin('username', TType.STRING, 2) - oprot.writeString(self.username.encode('utf-8') if sys.version_info[0] == 2 else self.username) - oprot.writeFieldEnd() - if self.password is not None: - oprot.writeFieldBegin('password', TType.STRING, 3) - oprot.writeString(self.password.encode('utf-8') if sys.version_info[0] == 2 else self.password) - oprot.writeFieldEnd() - if self.configuration is not None: - oprot.writeFieldBegin('configuration', TType.MAP, 4) - oprot.writeMapBegin(TType.STRING, TType.STRING, len(self.configuration)) - for kiter132, viter133 in self.configuration.items(): - oprot.writeString(kiter132.encode('utf-8') if sys.version_info[0] == 2 else kiter132) - oprot.writeString(viter133.encode('utf-8') if sys.version_info[0] == 2 else viter133) - oprot.writeMapEnd() - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.client_protocol is None: - raise TProtocolException(message='Required field client_protocol is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TOpenSessionResp(object): - """ - Attributes: - - status - - serverProtocolVersion - - sessionHandle - - configuration - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.I32, 'serverProtocolVersion', None, 9, ), # 2 - (3, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 3 - (4, TType.MAP, 'configuration', (TType.STRING, 'UTF8', TType.STRING, 'UTF8', False), None, ), # 4 - ) - - def __init__(self, status=None, serverProtocolVersion=thrift_spec[2][4], sessionHandle=None, configuration=None,): - self.status = status - self.serverProtocolVersion = serverProtocolVersion - self.sessionHandle = sessionHandle - self.configuration = configuration - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.I32: - self.serverProtocolVersion = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.MAP: - self.configuration = {} - (_ktype135, _vtype136, _size134) = iprot.readMapBegin() - for _i138 in range(_size134): - _key139 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - _val140 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - self.configuration[_key139] = _val140 - iprot.readMapEnd() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TOpenSessionResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.serverProtocolVersion is not None: - oprot.writeFieldBegin('serverProtocolVersion', TType.I32, 2) - oprot.writeI32(self.serverProtocolVersion) - oprot.writeFieldEnd() - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 3) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - if self.configuration is not None: - oprot.writeFieldBegin('configuration', TType.MAP, 4) - oprot.writeMapBegin(TType.STRING, TType.STRING, len(self.configuration)) - for kiter141, viter142 in self.configuration.items(): - oprot.writeString(kiter141.encode('utf-8') if sys.version_info[0] == 2 else kiter141) - oprot.writeString(viter142.encode('utf-8') if sys.version_info[0] == 2 else viter142) - oprot.writeMapEnd() - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - if self.serverProtocolVersion is None: - raise TProtocolException(message='Required field serverProtocolVersion is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TCloseSessionReq(object): - """ - Attributes: - - sessionHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - ) - - def __init__(self, sessionHandle=None,): - self.sessionHandle = sessionHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TCloseSessionReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TCloseSessionResp(object): - """ - Attributes: - - status - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - ) - - def __init__(self, status=None,): - self.status = status - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TCloseSessionResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetInfoValue(object): - """ - Attributes: - - stringValue - - smallIntValue - - integerBitmask - - integerFlag - - binaryValue - - lenValue - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRING, 'stringValue', 'UTF8', None, ), # 1 - (2, TType.I16, 'smallIntValue', None, None, ), # 2 - (3, TType.I32, 'integerBitmask', None, None, ), # 3 - (4, TType.I32, 'integerFlag', None, None, ), # 4 - (5, TType.I32, 'binaryValue', None, None, ), # 5 - (6, TType.I64, 'lenValue', None, None, ), # 6 - ) - - def __init__(self, stringValue=None, smallIntValue=None, integerBitmask=None, integerFlag=None, binaryValue=None, lenValue=None,): - self.stringValue = stringValue - self.smallIntValue = smallIntValue - self.integerBitmask = integerBitmask - self.integerFlag = integerFlag - self.binaryValue = binaryValue - self.lenValue = lenValue - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRING: - self.stringValue = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.I16: - self.smallIntValue = iprot.readI16() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.I32: - self.integerBitmask = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.I32: - self.integerFlag = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 5: - if ftype == TType.I32: - self.binaryValue = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 6: - if ftype == TType.I64: - self.lenValue = iprot.readI64() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetInfoValue') - if self.stringValue is not None: - oprot.writeFieldBegin('stringValue', TType.STRING, 1) - oprot.writeString(self.stringValue.encode('utf-8') if sys.version_info[0] == 2 else self.stringValue) - oprot.writeFieldEnd() - if self.smallIntValue is not None: - oprot.writeFieldBegin('smallIntValue', TType.I16, 2) - oprot.writeI16(self.smallIntValue) - oprot.writeFieldEnd() - if self.integerBitmask is not None: - oprot.writeFieldBegin('integerBitmask', TType.I32, 3) - oprot.writeI32(self.integerBitmask) - oprot.writeFieldEnd() - if self.integerFlag is not None: - oprot.writeFieldBegin('integerFlag', TType.I32, 4) - oprot.writeI32(self.integerFlag) - oprot.writeFieldEnd() - if self.binaryValue is not None: - oprot.writeFieldBegin('binaryValue', TType.I32, 5) - oprot.writeI32(self.binaryValue) - oprot.writeFieldEnd() - if self.lenValue is not None: - oprot.writeFieldBegin('lenValue', TType.I64, 6) - oprot.writeI64(self.lenValue) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetInfoReq(object): - """ - Attributes: - - sessionHandle - - infoType - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - (2, TType.I32, 'infoType', None, None, ), # 2 - ) - - def __init__(self, sessionHandle=None, infoType=None,): - self.sessionHandle = sessionHandle - self.infoType = infoType - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.I32: - self.infoType = iprot.readI32() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetInfoReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - if self.infoType is not None: - oprot.writeFieldBegin('infoType', TType.I32, 2) - oprot.writeI32(self.infoType) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - if self.infoType is None: - raise TProtocolException(message='Required field infoType is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetInfoResp(object): - """ - Attributes: - - status - - infoValue - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'infoValue', (TGetInfoValue, TGetInfoValue.thrift_spec), None, ), # 2 - ) - - def __init__(self, status=None, infoValue=None,): - self.status = status - self.infoValue = infoValue - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.infoValue = TGetInfoValue() - self.infoValue.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetInfoResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.infoValue is not None: - oprot.writeFieldBegin('infoValue', TType.STRUCT, 2) - self.infoValue.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - if self.infoValue is None: - raise TProtocolException(message='Required field infoValue is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TExecuteStatementReq(object): - """ - Attributes: - - sessionHandle - - statement - - confOverlay - - runAsync - - queryTimeout - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - (2, TType.STRING, 'statement', 'UTF8', None, ), # 2 - (3, TType.MAP, 'confOverlay', (TType.STRING, 'UTF8', TType.STRING, 'UTF8', False), None, ), # 3 - (4, TType.BOOL, 'runAsync', None, False, ), # 4 - (5, TType.I64, 'queryTimeout', None, 0, ), # 5 - ) - - def __init__(self, sessionHandle=None, statement=None, confOverlay=None, runAsync=thrift_spec[4][4], queryTimeout=thrift_spec[5][4],): - self.sessionHandle = sessionHandle - self.statement = statement - self.confOverlay = confOverlay - self.runAsync = runAsync - self.queryTimeout = queryTimeout - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.statement = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.MAP: - self.confOverlay = {} - (_ktype144, _vtype145, _size143) = iprot.readMapBegin() - for _i147 in range(_size143): - _key148 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - _val149 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - self.confOverlay[_key148] = _val149 - iprot.readMapEnd() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.BOOL: - self.runAsync = iprot.readBool() - else: - iprot.skip(ftype) - elif fid == 5: - if ftype == TType.I64: - self.queryTimeout = iprot.readI64() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TExecuteStatementReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - if self.statement is not None: - oprot.writeFieldBegin('statement', TType.STRING, 2) - oprot.writeString(self.statement.encode('utf-8') if sys.version_info[0] == 2 else self.statement) - oprot.writeFieldEnd() - if self.confOverlay is not None: - oprot.writeFieldBegin('confOverlay', TType.MAP, 3) - oprot.writeMapBegin(TType.STRING, TType.STRING, len(self.confOverlay)) - for kiter150, viter151 in self.confOverlay.items(): - oprot.writeString(kiter150.encode('utf-8') if sys.version_info[0] == 2 else kiter150) - oprot.writeString(viter151.encode('utf-8') if sys.version_info[0] == 2 else viter151) - oprot.writeMapEnd() - oprot.writeFieldEnd() - if self.runAsync is not None: - oprot.writeFieldBegin('runAsync', TType.BOOL, 4) - oprot.writeBool(self.runAsync) - oprot.writeFieldEnd() - if self.queryTimeout is not None: - oprot.writeFieldBegin('queryTimeout', TType.I64, 5) - oprot.writeI64(self.queryTimeout) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - if self.statement is None: - raise TProtocolException(message='Required field statement is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TExecuteStatementResp(object): - """ - Attributes: - - status - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 2 - ) - - def __init__(self, status=None, operationHandle=None,): - self.status = status - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TExecuteStatementResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 2) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetTypeInfoReq(object): - """ - Attributes: - - sessionHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - ) - - def __init__(self, sessionHandle=None,): - self.sessionHandle = sessionHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetTypeInfoReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetTypeInfoResp(object): - """ - Attributes: - - status - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 2 - ) - - def __init__(self, status=None, operationHandle=None,): - self.status = status - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetTypeInfoResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 2) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetCatalogsReq(object): - """ - Attributes: - - sessionHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - ) - - def __init__(self, sessionHandle=None,): - self.sessionHandle = sessionHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetCatalogsReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetCatalogsResp(object): - """ - Attributes: - - status - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 2 - ) - - def __init__(self, status=None, operationHandle=None,): - self.status = status - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetCatalogsResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 2) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetSchemasReq(object): - """ - Attributes: - - sessionHandle - - catalogName - - schemaName - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - (2, TType.STRING, 'catalogName', 'UTF8', None, ), # 2 - (3, TType.STRING, 'schemaName', 'UTF8', None, ), # 3 - ) - - def __init__(self, sessionHandle=None, catalogName=None, schemaName=None,): - self.sessionHandle = sessionHandle - self.catalogName = catalogName - self.schemaName = schemaName - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.catalogName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRING: - self.schemaName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetSchemasReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - if self.catalogName is not None: - oprot.writeFieldBegin('catalogName', TType.STRING, 2) - oprot.writeString(self.catalogName.encode('utf-8') if sys.version_info[0] == 2 else self.catalogName) - oprot.writeFieldEnd() - if self.schemaName is not None: - oprot.writeFieldBegin('schemaName', TType.STRING, 3) - oprot.writeString(self.schemaName.encode('utf-8') if sys.version_info[0] == 2 else self.schemaName) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetSchemasResp(object): - """ - Attributes: - - status - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 2 - ) - - def __init__(self, status=None, operationHandle=None,): - self.status = status - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetSchemasResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 2) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetTablesReq(object): - """ - Attributes: - - sessionHandle - - catalogName - - schemaName - - tableName - - tableTypes - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - (2, TType.STRING, 'catalogName', 'UTF8', None, ), # 2 - (3, TType.STRING, 'schemaName', 'UTF8', None, ), # 3 - (4, TType.STRING, 'tableName', 'UTF8', None, ), # 4 - (5, TType.LIST, 'tableTypes', (TType.STRING, 'UTF8', False), None, ), # 5 - ) - - def __init__(self, sessionHandle=None, catalogName=None, schemaName=None, tableName=None, tableTypes=None,): - self.sessionHandle = sessionHandle - self.catalogName = catalogName - self.schemaName = schemaName - self.tableName = tableName - self.tableTypes = tableTypes - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.catalogName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRING: - self.schemaName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.STRING: - self.tableName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 5: - if ftype == TType.LIST: - self.tableTypes = [] - (_etype155, _size152) = iprot.readListBegin() - for _i156 in range(_size152): - _elem157 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - self.tableTypes.append(_elem157) - iprot.readListEnd() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetTablesReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - if self.catalogName is not None: - oprot.writeFieldBegin('catalogName', TType.STRING, 2) - oprot.writeString(self.catalogName.encode('utf-8') if sys.version_info[0] == 2 else self.catalogName) - oprot.writeFieldEnd() - if self.schemaName is not None: - oprot.writeFieldBegin('schemaName', TType.STRING, 3) - oprot.writeString(self.schemaName.encode('utf-8') if sys.version_info[0] == 2 else self.schemaName) - oprot.writeFieldEnd() - if self.tableName is not None: - oprot.writeFieldBegin('tableName', TType.STRING, 4) - oprot.writeString(self.tableName.encode('utf-8') if sys.version_info[0] == 2 else self.tableName) - oprot.writeFieldEnd() - if self.tableTypes is not None: - oprot.writeFieldBegin('tableTypes', TType.LIST, 5) - oprot.writeListBegin(TType.STRING, len(self.tableTypes)) - for iter158 in self.tableTypes: - oprot.writeString(iter158.encode('utf-8') if sys.version_info[0] == 2 else iter158) - oprot.writeListEnd() - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetTablesResp(object): - """ - Attributes: - - status - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 2 - ) - - def __init__(self, status=None, operationHandle=None,): - self.status = status - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetTablesResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 2) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetTableTypesReq(object): - """ - Attributes: - - sessionHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - ) - - def __init__(self, sessionHandle=None,): - self.sessionHandle = sessionHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetTableTypesReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetTableTypesResp(object): - """ - Attributes: - - status - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 2 - ) - - def __init__(self, status=None, operationHandle=None,): - self.status = status - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetTableTypesResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 2) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetColumnsReq(object): - """ - Attributes: - - sessionHandle - - catalogName - - schemaName - - tableName - - columnName - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - (2, TType.STRING, 'catalogName', 'UTF8', None, ), # 2 - (3, TType.STRING, 'schemaName', 'UTF8', None, ), # 3 - (4, TType.STRING, 'tableName', 'UTF8', None, ), # 4 - (5, TType.STRING, 'columnName', 'UTF8', None, ), # 5 - ) - - def __init__(self, sessionHandle=None, catalogName=None, schemaName=None, tableName=None, columnName=None,): - self.sessionHandle = sessionHandle - self.catalogName = catalogName - self.schemaName = schemaName - self.tableName = tableName - self.columnName = columnName - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.catalogName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRING: - self.schemaName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.STRING: - self.tableName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 5: - if ftype == TType.STRING: - self.columnName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetColumnsReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - if self.catalogName is not None: - oprot.writeFieldBegin('catalogName', TType.STRING, 2) - oprot.writeString(self.catalogName.encode('utf-8') if sys.version_info[0] == 2 else self.catalogName) - oprot.writeFieldEnd() - if self.schemaName is not None: - oprot.writeFieldBegin('schemaName', TType.STRING, 3) - oprot.writeString(self.schemaName.encode('utf-8') if sys.version_info[0] == 2 else self.schemaName) - oprot.writeFieldEnd() - if self.tableName is not None: - oprot.writeFieldBegin('tableName', TType.STRING, 4) - oprot.writeString(self.tableName.encode('utf-8') if sys.version_info[0] == 2 else self.tableName) - oprot.writeFieldEnd() - if self.columnName is not None: - oprot.writeFieldBegin('columnName', TType.STRING, 5) - oprot.writeString(self.columnName.encode('utf-8') if sys.version_info[0] == 2 else self.columnName) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetColumnsResp(object): - """ - Attributes: - - status - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 2 - ) - - def __init__(self, status=None, operationHandle=None,): - self.status = status - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetColumnsResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 2) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetFunctionsReq(object): - """ - Attributes: - - sessionHandle - - catalogName - - schemaName - - functionName - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - (2, TType.STRING, 'catalogName', 'UTF8', None, ), # 2 - (3, TType.STRING, 'schemaName', 'UTF8', None, ), # 3 - (4, TType.STRING, 'functionName', 'UTF8', None, ), # 4 - ) - - def __init__(self, sessionHandle=None, catalogName=None, schemaName=None, functionName=None,): - self.sessionHandle = sessionHandle - self.catalogName = catalogName - self.schemaName = schemaName - self.functionName = functionName - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.catalogName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRING: - self.schemaName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.STRING: - self.functionName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetFunctionsReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - if self.catalogName is not None: - oprot.writeFieldBegin('catalogName', TType.STRING, 2) - oprot.writeString(self.catalogName.encode('utf-8') if sys.version_info[0] == 2 else self.catalogName) - oprot.writeFieldEnd() - if self.schemaName is not None: - oprot.writeFieldBegin('schemaName', TType.STRING, 3) - oprot.writeString(self.schemaName.encode('utf-8') if sys.version_info[0] == 2 else self.schemaName) - oprot.writeFieldEnd() - if self.functionName is not None: - oprot.writeFieldBegin('functionName', TType.STRING, 4) - oprot.writeString(self.functionName.encode('utf-8') if sys.version_info[0] == 2 else self.functionName) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - if self.functionName is None: - raise TProtocolException(message='Required field functionName is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetFunctionsResp(object): - """ - Attributes: - - status - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 2 - ) - - def __init__(self, status=None, operationHandle=None,): - self.status = status - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetFunctionsResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 2) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetPrimaryKeysReq(object): - """ - Attributes: - - sessionHandle - - catalogName - - schemaName - - tableName - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - (2, TType.STRING, 'catalogName', 'UTF8', None, ), # 2 - (3, TType.STRING, 'schemaName', 'UTF8', None, ), # 3 - (4, TType.STRING, 'tableName', 'UTF8', None, ), # 4 - ) - - def __init__(self, sessionHandle=None, catalogName=None, schemaName=None, tableName=None,): - self.sessionHandle = sessionHandle - self.catalogName = catalogName - self.schemaName = schemaName - self.tableName = tableName - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.catalogName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRING: - self.schemaName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.STRING: - self.tableName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetPrimaryKeysReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - if self.catalogName is not None: - oprot.writeFieldBegin('catalogName', TType.STRING, 2) - oprot.writeString(self.catalogName.encode('utf-8') if sys.version_info[0] == 2 else self.catalogName) - oprot.writeFieldEnd() - if self.schemaName is not None: - oprot.writeFieldBegin('schemaName', TType.STRING, 3) - oprot.writeString(self.schemaName.encode('utf-8') if sys.version_info[0] == 2 else self.schemaName) - oprot.writeFieldEnd() - if self.tableName is not None: - oprot.writeFieldBegin('tableName', TType.STRING, 4) - oprot.writeString(self.tableName.encode('utf-8') if sys.version_info[0] == 2 else self.tableName) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetPrimaryKeysResp(object): - """ - Attributes: - - status - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 2 - ) - - def __init__(self, status=None, operationHandle=None,): - self.status = status - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetPrimaryKeysResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 2) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetCrossReferenceReq(object): - """ - Attributes: - - sessionHandle - - parentCatalogName - - parentSchemaName - - parentTableName - - foreignCatalogName - - foreignSchemaName - - foreignTableName - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - (2, TType.STRING, 'parentCatalogName', 'UTF8', None, ), # 2 - (3, TType.STRING, 'parentSchemaName', 'UTF8', None, ), # 3 - (4, TType.STRING, 'parentTableName', 'UTF8', None, ), # 4 - (5, TType.STRING, 'foreignCatalogName', 'UTF8', None, ), # 5 - (6, TType.STRING, 'foreignSchemaName', 'UTF8', None, ), # 6 - (7, TType.STRING, 'foreignTableName', 'UTF8', None, ), # 7 - ) - - def __init__(self, sessionHandle=None, parentCatalogName=None, parentSchemaName=None, parentTableName=None, foreignCatalogName=None, foreignSchemaName=None, foreignTableName=None,): - self.sessionHandle = sessionHandle - self.parentCatalogName = parentCatalogName - self.parentSchemaName = parentSchemaName - self.parentTableName = parentTableName - self.foreignCatalogName = foreignCatalogName - self.foreignSchemaName = foreignSchemaName - self.foreignTableName = foreignTableName - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.parentCatalogName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRING: - self.parentSchemaName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.STRING: - self.parentTableName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 5: - if ftype == TType.STRING: - self.foreignCatalogName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 6: - if ftype == TType.STRING: - self.foreignSchemaName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 7: - if ftype == TType.STRING: - self.foreignTableName = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetCrossReferenceReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - if self.parentCatalogName is not None: - oprot.writeFieldBegin('parentCatalogName', TType.STRING, 2) - oprot.writeString(self.parentCatalogName.encode('utf-8') if sys.version_info[0] == 2 else self.parentCatalogName) - oprot.writeFieldEnd() - if self.parentSchemaName is not None: - oprot.writeFieldBegin('parentSchemaName', TType.STRING, 3) - oprot.writeString(self.parentSchemaName.encode('utf-8') if sys.version_info[0] == 2 else self.parentSchemaName) - oprot.writeFieldEnd() - if self.parentTableName is not None: - oprot.writeFieldBegin('parentTableName', TType.STRING, 4) - oprot.writeString(self.parentTableName.encode('utf-8') if sys.version_info[0] == 2 else self.parentTableName) - oprot.writeFieldEnd() - if self.foreignCatalogName is not None: - oprot.writeFieldBegin('foreignCatalogName', TType.STRING, 5) - oprot.writeString(self.foreignCatalogName.encode('utf-8') if sys.version_info[0] == 2 else self.foreignCatalogName) - oprot.writeFieldEnd() - if self.foreignSchemaName is not None: - oprot.writeFieldBegin('foreignSchemaName', TType.STRING, 6) - oprot.writeString(self.foreignSchemaName.encode('utf-8') if sys.version_info[0] == 2 else self.foreignSchemaName) - oprot.writeFieldEnd() - if self.foreignTableName is not None: - oprot.writeFieldBegin('foreignTableName', TType.STRING, 7) - oprot.writeString(self.foreignTableName.encode('utf-8') if sys.version_info[0] == 2 else self.foreignTableName) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetCrossReferenceResp(object): - """ - Attributes: - - status - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 2 - ) - - def __init__(self, status=None, operationHandle=None,): - self.status = status - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetCrossReferenceResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 2) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - -class TProgressUpdateResp(object): - """ - Attributes: - - headerNames - - rows - - progressedPercentage - - status - - footerSummary - - startTime - """ - - thrift_spec = ( - None, # 0 - (1, TType.LIST, 'headerNames', (TType.STRING, 'UTF8', False), None, ), # 1 - (2, TType.LIST, 'rows', (TType.LIST, (TType.STRING, 'UTF8', False), False), None, ), # 2 - (3, TType.DOUBLE, 'progressedPercentage', None, None, ), # 3 - (4, TType.I32, 'status', None, None, ), # 4 - (5, TType.STRING, 'footerSummary', 'UTF8', None, ), # 5 - (6, TType.I64, 'startTime', None, None, ), # 6 - ) - - def __init__(self, headerNames=None, rows=None, progressedPercentage=None, status=None, footerSummary=None, startTime=None,): - self.headerNames = headerNames - self.rows = rows - self.progressedPercentage = progressedPercentage - self.status = status - self.footerSummary = footerSummary - self.startTime = startTime - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.LIST: - self.headerNames = [] - (_etype162, _size159) = iprot.readListBegin() - for _i163 in range(_size159): - _elem164 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - self.headerNames.append(_elem164) - iprot.readListEnd() - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.LIST: - self.rows = [] - (_etype168, _size165) = iprot.readListBegin() - for _i169 in range(_size165): - _elem170 = [] - (_etype174, _size171) = iprot.readListBegin() - for _i175 in range(_size171): - _elem176 = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - _elem170.append(_elem176) - iprot.readListEnd() - self.rows.append(_elem170) - iprot.readListEnd() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.DOUBLE: - self.progressedPercentage = iprot.readDouble() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.I32: - self.status = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 5: - if ftype == TType.STRING: - self.footerSummary = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 6: - if ftype == TType.I64: - self.startTime = iprot.readI64() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TProgressUpdateResp') - if self.headerNames is not None: - oprot.writeFieldBegin('headerNames', TType.LIST, 1) - oprot.writeListBegin(TType.STRING, len(self.headerNames)) - for iter177 in self.headerNames: - oprot.writeString(iter177.encode('utf-8') if sys.version_info[0] == 2 else iter177) - oprot.writeListEnd() - oprot.writeFieldEnd() - if self.rows is not None: - oprot.writeFieldBegin('rows', TType.LIST, 2) - oprot.writeListBegin(TType.LIST, len(self.rows)) - for iter178 in self.rows: - oprot.writeListBegin(TType.STRING, len(iter178)) - for iter179 in iter178: - oprot.writeString(iter179.encode('utf-8') if sys.version_info[0] == 2 else iter179) - oprot.writeListEnd() - oprot.writeListEnd() - oprot.writeFieldEnd() - if self.progressedPercentage is not None: - oprot.writeFieldBegin('progressedPercentage', TType.DOUBLE, 3) - oprot.writeDouble(self.progressedPercentage) - oprot.writeFieldEnd() - if self.status is not None: - oprot.writeFieldBegin('status', TType.I32, 4) - oprot.writeI32(self.status) - oprot.writeFieldEnd() - if self.footerSummary is not None: - oprot.writeFieldBegin('footerSummary', TType.STRING, 5) - oprot.writeString(self.footerSummary.encode('utf-8') if sys.version_info[0] == 2 else self.footerSummary) - oprot.writeFieldEnd() - if self.startTime is not None: - oprot.writeFieldBegin('startTime', TType.I64, 6) - oprot.writeI64(self.startTime) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.headerNames is None: - raise TProtocolException(message='Required field headerNames is unset!') - if self.rows is None: - raise TProtocolException(message='Required field rows is unset!') - if self.progressedPercentage is None: - raise TProtocolException(message='Required field progressedPercentage is unset!') - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - if self.footerSummary is None: - raise TProtocolException(message='Required field footerSummary is unset!') - if self.startTime is None: - raise TProtocolException(message='Required field startTime is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - -class TGetOperationStatusReq(object): - """ - Attributes: - - operationHandle - - getProgressUpdate - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 1 - (2, TType.BOOL, 'getProgressUpdate', None, None, ), # 2 - ) - - def __init__(self, operationHandle=None, getProgressUpdate=None,): - self.operationHandle = operationHandle - self.getProgressUpdate = getProgressUpdate - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.BOOL: - self.getProgressUpdate = iprot.readBool() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetOperationStatusReq') - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 1) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - if self.getProgressUpdate is not None: - oprot.writeFieldBegin('getProgressUpdate', TType.BOOL, 2) - oprot.writeBool(self.getProgressUpdate) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.operationHandle is None: - raise TProtocolException(message='Required field operationHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetOperationStatusResp(object): - """ - Attributes: - - status - - operationState - - sqlState - - errorCode - - errorMessage - - taskStatus - - operationStarted - - operationCompleted - - hasResultSet - - progressUpdateResponse - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.I32, 'operationState', None, None, ), # 2 - (3, TType.STRING, 'sqlState', 'UTF8', None, ), # 3 - (4, TType.I32, 'errorCode', None, None, ), # 4 - (5, TType.STRING, 'errorMessage', 'UTF8', None, ), # 5 - (6, TType.STRING, 'taskStatus', 'UTF8', None, ), # 6 - (7, TType.I64, 'operationStarted', None, None, ), # 7 - (8, TType.I64, 'operationCompleted', None, None, ), # 8 - (9, TType.BOOL, 'hasResultSet', None, None, ), # 9 - (10, TType.STRUCT, 'progressUpdateResponse', (TProgressUpdateResp, TProgressUpdateResp.thrift_spec), None, ), # 10 - ) - - def __init__(self, status=None, operationState=None, sqlState=None, errorCode=None, errorMessage=None, taskStatus=None, operationStarted=None, operationCompleted=None, hasResultSet=None, progressUpdateResponse=None,): - self.status = status - self.operationState = operationState - self.sqlState = sqlState - self.errorCode = errorCode - self.errorMessage = errorMessage - self.taskStatus = taskStatus - self.operationStarted = operationStarted - self.operationCompleted = operationCompleted - self.hasResultSet = hasResultSet - self.progressUpdateResponse = progressUpdateResponse - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.I32: - self.operationState = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRING: - self.sqlState = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.I32: - self.errorCode = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 5: - if ftype == TType.STRING: - self.errorMessage = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 6: - if ftype == TType.STRING: - self.taskStatus = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 7: - if ftype == TType.I64: - self.operationStarted = iprot.readI64() - else: - iprot.skip(ftype) - elif fid == 8: - if ftype == TType.I64: - self.operationCompleted = iprot.readI64() - else: - iprot.skip(ftype) - elif fid == 9: - if ftype == TType.BOOL: - self.hasResultSet = iprot.readBool() - else: - iprot.skip(ftype) - elif fid == 10: - if ftype == TType.STRUCT: - self.progressUpdateResponse = TProgressUpdateResp() - self.progressUpdateResponse.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetOperationStatusResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.operationState is not None: - oprot.writeFieldBegin('operationState', TType.I32, 2) - oprot.writeI32(self.operationState) - oprot.writeFieldEnd() - if self.sqlState is not None: - oprot.writeFieldBegin('sqlState', TType.STRING, 3) - oprot.writeString(self.sqlState.encode('utf-8') if sys.version_info[0] == 2 else self.sqlState) - oprot.writeFieldEnd() - if self.errorCode is not None: - oprot.writeFieldBegin('errorCode', TType.I32, 4) - oprot.writeI32(self.errorCode) - oprot.writeFieldEnd() - if self.errorMessage is not None: - oprot.writeFieldBegin('errorMessage', TType.STRING, 5) - oprot.writeString(self.errorMessage.encode('utf-8') if sys.version_info[0] == 2 else self.errorMessage) - oprot.writeFieldEnd() - if self.taskStatus is not None: - oprot.writeFieldBegin('taskStatus', TType.STRING, 6) - oprot.writeString(self.taskStatus.encode('utf-8') if sys.version_info[0] == 2 else self.taskStatus) - oprot.writeFieldEnd() - if self.operationStarted is not None: - oprot.writeFieldBegin('operationStarted', TType.I64, 7) - oprot.writeI64(self.operationStarted) - oprot.writeFieldEnd() - if self.operationCompleted is not None: - oprot.writeFieldBegin('operationCompleted', TType.I64, 8) - oprot.writeI64(self.operationCompleted) - oprot.writeFieldEnd() - if self.hasResultSet is not None: - oprot.writeFieldBegin('hasResultSet', TType.BOOL, 9) - oprot.writeBool(self.hasResultSet) - oprot.writeFieldEnd() - if self.progressUpdateResponse is not None: - oprot.writeFieldBegin('progressUpdateResponse', TType.STRUCT, 10) - self.progressUpdateResponse.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TCancelOperationReq(object): - """ - Attributes: - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 1 - ) - - def __init__(self, operationHandle=None,): - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TCancelOperationReq') - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 1) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.operationHandle is None: - raise TProtocolException(message='Required field operationHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TCancelOperationResp(object): - """ - Attributes: - - status - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - ) - - def __init__(self, status=None,): - self.status = status - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TCancelOperationResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TCloseOperationReq(object): - """ - Attributes: - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 1 - ) - - def __init__(self, operationHandle=None,): - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TCloseOperationReq') - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 1) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.operationHandle is None: - raise TProtocolException(message='Required field operationHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TCloseOperationResp(object): - """ - Attributes: - - status - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - ) - - def __init__(self, status=None,): - self.status = status - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TCloseOperationResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetResultSetMetadataReq(object): - """ - Attributes: - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 1 - ) - - def __init__(self, operationHandle=None,): - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetResultSetMetadataReq') - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 1) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.operationHandle is None: - raise TProtocolException(message='Required field operationHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetResultSetMetadataResp(object): - """ - Attributes: - - status - - schema - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRUCT, 'schema', (TTableSchema, TTableSchema.thrift_spec), None, ), # 2 - ) - - def __init__(self, status=None, schema=None,): - self.status = status - self.schema = schema - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRUCT: - self.schema = TTableSchema() - self.schema.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetResultSetMetadataResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.schema is not None: - oprot.writeFieldBegin('schema', TType.STRUCT, 2) - self.schema.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TFetchResultsReq(object): - """ - Attributes: - - operationHandle - - orientation - - maxRows - - fetchType - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 1 - (2, TType.I32, 'orientation', None, 0, ), # 2 - (3, TType.I64, 'maxRows', None, None, ), # 3 - (4, TType.I16, 'fetchType', None, 0, ), # 4 - ) - - def __init__(self, operationHandle=None, orientation=thrift_spec[2][4], maxRows=None, fetchType=thrift_spec[4][4],): - self.operationHandle = operationHandle - self.orientation = orientation - self.maxRows = maxRows - self.fetchType = fetchType - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.I32: - self.orientation = iprot.readI32() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.I64: - self.maxRows = iprot.readI64() - else: - iprot.skip(ftype) - elif fid == 4: - if ftype == TType.I16: - self.fetchType = iprot.readI16() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TFetchResultsReq') - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 1) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - if self.orientation is not None: - oprot.writeFieldBegin('orientation', TType.I32, 2) - oprot.writeI32(self.orientation) - oprot.writeFieldEnd() - if self.maxRows is not None: - oprot.writeFieldBegin('maxRows', TType.I64, 3) - oprot.writeI64(self.maxRows) - oprot.writeFieldEnd() - if self.fetchType is not None: - oprot.writeFieldBegin('fetchType', TType.I16, 4) - oprot.writeI16(self.fetchType) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.operationHandle is None: - raise TProtocolException(message='Required field operationHandle is unset!') - if self.orientation is None: - raise TProtocolException(message='Required field orientation is unset!') - if self.maxRows is None: - raise TProtocolException(message='Required field maxRows is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TFetchResultsResp(object): - """ - Attributes: - - status - - hasMoreRows - - results - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.BOOL, 'hasMoreRows', None, None, ), # 2 - (3, TType.STRUCT, 'results', (TRowSet, TRowSet.thrift_spec), None, ), # 3 - ) - - def __init__(self, status=None, hasMoreRows=None, results=None,): - self.status = status - self.hasMoreRows = hasMoreRows - self.results = results - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.BOOL: - self.hasMoreRows = iprot.readBool() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRUCT: - self.results = TRowSet() - self.results.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TFetchResultsResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.hasMoreRows is not None: - oprot.writeFieldBegin('hasMoreRows', TType.BOOL, 2) - oprot.writeBool(self.hasMoreRows) - oprot.writeFieldEnd() - if self.results is not None: - oprot.writeFieldBegin('results', TType.STRUCT, 3) - self.results.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetDelegationTokenReq(object): - """ - Attributes: - - sessionHandle - - owner - - renewer - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - (2, TType.STRING, 'owner', 'UTF8', None, ), # 2 - (3, TType.STRING, 'renewer', 'UTF8', None, ), # 3 - ) - - def __init__(self, sessionHandle=None, owner=None, renewer=None,): - self.sessionHandle = sessionHandle - self.owner = owner - self.renewer = renewer - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.owner = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - elif fid == 3: - if ftype == TType.STRING: - self.renewer = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetDelegationTokenReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - if self.owner is not None: - oprot.writeFieldBegin('owner', TType.STRING, 2) - oprot.writeString(self.owner.encode('utf-8') if sys.version_info[0] == 2 else self.owner) - oprot.writeFieldEnd() - if self.renewer is not None: - oprot.writeFieldBegin('renewer', TType.STRING, 3) - oprot.writeString(self.renewer.encode('utf-8') if sys.version_info[0] == 2 else self.renewer) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - if self.owner is None: - raise TProtocolException(message='Required field owner is unset!') - if self.renewer is None: - raise TProtocolException(message='Required field renewer is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetDelegationTokenResp(object): - """ - Attributes: - - status - - delegationToken - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRING, 'delegationToken', 'UTF8', None, ), # 2 - ) - - def __init__(self, status=None, delegationToken=None,): - self.status = status - self.delegationToken = delegationToken - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.delegationToken = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetDelegationTokenResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.delegationToken is not None: - oprot.writeFieldBegin('delegationToken', TType.STRING, 2) - oprot.writeString(self.delegationToken.encode('utf-8') if sys.version_info[0] == 2 else self.delegationToken) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TCancelDelegationTokenReq(object): - """ - Attributes: - - sessionHandle - - delegationToken - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - (2, TType.STRING, 'delegationToken', 'UTF8', None, ), # 2 - ) - - def __init__(self, sessionHandle=None, delegationToken=None,): - self.sessionHandle = sessionHandle - self.delegationToken = delegationToken - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.delegationToken = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TCancelDelegationTokenReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - if self.delegationToken is not None: - oprot.writeFieldBegin('delegationToken', TType.STRING, 2) - oprot.writeString(self.delegationToken.encode('utf-8') if sys.version_info[0] == 2 else self.delegationToken) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - if self.delegationToken is None: - raise TProtocolException(message='Required field delegationToken is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TCancelDelegationTokenResp(object): - """ - Attributes: - - status - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - ) - - def __init__(self, status=None,): - self.status = status - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TCancelDelegationTokenResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TRenewDelegationTokenReq(object): - """ - Attributes: - - sessionHandle - - delegationToken - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'sessionHandle', (TSessionHandle, TSessionHandle.thrift_spec), None, ), # 1 - (2, TType.STRING, 'delegationToken', 'UTF8', None, ), # 2 - ) - - def __init__(self, sessionHandle=None, delegationToken=None,): - self.sessionHandle = sessionHandle - self.delegationToken = delegationToken - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.sessionHandle = TSessionHandle() - self.sessionHandle.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.delegationToken = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TRenewDelegationTokenReq') - if self.sessionHandle is not None: - oprot.writeFieldBegin('sessionHandle', TType.STRUCT, 1) - self.sessionHandle.write(oprot) - oprot.writeFieldEnd() - if self.delegationToken is not None: - oprot.writeFieldBegin('delegationToken', TType.STRING, 2) - oprot.writeString(self.delegationToken.encode('utf-8') if sys.version_info[0] == 2 else self.delegationToken) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.sessionHandle is None: - raise TProtocolException(message='Required field sessionHandle is unset!') - if self.delegationToken is None: - raise TProtocolException(message='Required field delegationToken is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TRenewDelegationTokenResp(object): - """ - Attributes: - - status - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - ) - - def __init__(self, status=None,): - self.status = status - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TRenewDelegationTokenResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - -class TGetLogReq(object): - """ - Attributes: - - operationHandle - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'operationHandle', (TOperationHandle, TOperationHandle.thrift_spec), None, ), # 1 - ) - - def __init__(self, operationHandle=None,): - self.operationHandle = operationHandle - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.operationHandle = TOperationHandle() - self.operationHandle.read(iprot) - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetLogReq') - if self.operationHandle is not None: - oprot.writeFieldBegin('operationHandle', TType.STRUCT, 1) - self.operationHandle.write(oprot) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.operationHandle is None: - raise TProtocolException(message='Required field operationHandle is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) - - -class TGetLogResp(object): - """ - Attributes: - - status - - log - """ - - thrift_spec = ( - None, # 0 - (1, TType.STRUCT, 'status', (TStatus, TStatus.thrift_spec), None, ), # 1 - (2, TType.STRING, 'log', 'UTF8', None, ), # 2 - ) - - def __init__(self, status=None, log=None,): - self.status = status - self.log = log - - def read(self, iprot): - if iprot._fast_decode is not None and isinstance(iprot.trans, TTransport.CReadableTransport) and self.thrift_spec is not None: - iprot._fast_decode(self, iprot, (self.__class__, self.thrift_spec)) - return - iprot.readStructBegin() - while True: - (fname, ftype, fid) = iprot.readFieldBegin() - if ftype == TType.STOP: - break - if fid == 1: - if ftype == TType.STRUCT: - self.status = TStatus() - self.status.read(iprot) - else: - iprot.skip(ftype) - elif fid == 2: - if ftype == TType.STRING: - self.log = iprot.readString().decode('utf-8') if sys.version_info[0] == 2 else iprot.readString() - else: - iprot.skip(ftype) - else: - iprot.skip(ftype) - iprot.readFieldEnd() - iprot.readStructEnd() - - def write(self, oprot): - if oprot._fast_encode is not None and self.thrift_spec is not None: - oprot.trans.write(oprot._fast_encode(self, (self.__class__, self.thrift_spec))) - return - oprot.writeStructBegin('TGetLogResp') - if self.status is not None: - oprot.writeFieldBegin('status', TType.STRUCT, 1) - self.status.write(oprot) - oprot.writeFieldEnd() - if self.log is not None: - oprot.writeFieldBegin('log', TType.STRING, 2) - oprot.writeString(self.log.encode('utf-8') if sys.version_info[0] == 2 else self.log) - oprot.writeFieldEnd() - oprot.writeFieldStop() - oprot.writeStructEnd() - - def validate(self): - if self.status is None: - raise TProtocolException(message='Required field status is unset!') - if self.log is None: - raise TProtocolException(message='Required field log is unset!') - return - - def __repr__(self): - L = ['%s=%r' % (key, value) - for key, value in self.__dict__.items()] - return '%s(%s)' % (self.__class__.__name__, ', '.join(L)) - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - - def __ne__(self, other): - return not (self == other) diff --git a/src/chronify/_vendor/kyuubi/pyhive/__init__.py b/src/chronify/_vendor/kyuubi/pyhive/__init__.py deleted file mode 100644 index 0a6bb1f..0000000 --- a/src/chronify/_vendor/kyuubi/pyhive/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from __future__ import absolute_import -from __future__ import unicode_literals -__version__ = '0.7.0' diff --git a/src/chronify/_vendor/kyuubi/pyhive/common.py b/src/chronify/_vendor/kyuubi/pyhive/common.py deleted file mode 100644 index 51692b9..0000000 --- a/src/chronify/_vendor/kyuubi/pyhive/common.py +++ /dev/null @@ -1,266 +0,0 @@ -"""Package private common utilities. Do not use directly. - -Many docstrings in this file are based on PEP-249, which is in the public domain. -""" - -from __future__ import absolute_import -from __future__ import unicode_literals -from builtins import bytes -from builtins import int -from builtins import object -from builtins import str -from past.builtins import basestring -from pyhive import exc -import abc -import collections -import time -import datetime -from future.utils import with_metaclass -from itertools import islice - -try: - from collections.abc import Iterable -except ImportError: - from collections import Iterable - - -class DBAPICursor(with_metaclass(abc.ABCMeta, object)): - """Base class for some common DB-API logic""" - - _STATE_NONE = 0 - _STATE_RUNNING = 1 - _STATE_FINISHED = 2 - - def __init__(self, poll_interval=1): - self._poll_interval = poll_interval - self._reset_state() - self.lastrowid = None - - def _reset_state(self): - """Reset state about the previous query in preparation for running another query""" - # State to return as part of DB-API - self._rownumber = 0 - - # Internal helper state - self._state = self._STATE_NONE - self._data = collections.deque() - self._columns = None - - def _fetch_while(self, fn): - while fn(): - self._fetch_more() - if fn(): - time.sleep(self._poll_interval) - - @abc.abstractproperty - def description(self): - raise NotImplementedError # pragma: no cover - - def close(self): - """By default, do nothing""" - pass - - @abc.abstractmethod - def _fetch_more(self): - """Get more results, append it to ``self._data``, and update ``self._state``.""" - raise NotImplementedError # pragma: no cover - - @property - def rowcount(self): - """By default, return -1 to indicate that this is not supported.""" - return -1 - - @abc.abstractmethod - def execute(self, operation, parameters=None): - """Prepare and execute a database operation (query or command). - - Parameters may be provided as sequence or mapping and will be bound to variables in the - operation. Variables are specified in a database-specific notation (see the module's - ``paramstyle`` attribute for details). - - Return values are not defined. - """ - raise NotImplementedError # pragma: no cover - - def executemany(self, operation, seq_of_parameters): - """Prepare a database operation (query or command) and then execute it against all parameter - sequences or mappings found in the sequence ``seq_of_parameters``. - - Only the final result set is retained. - - Return values are not defined. - """ - for parameters in seq_of_parameters[:-1]: - self.execute(operation, parameters) - while self._state != self._STATE_FINISHED: - self._fetch_more() - if seq_of_parameters: - self.execute(operation, seq_of_parameters[-1]) - - def fetchone(self): - """Fetch the next row of a query result set, returning a single sequence, or ``None`` when - no more data is available. - - An :py:class:`~pyhive.exc.Error` (or subclass) exception is raised if the previous call to - :py:meth:`execute` did not produce any result set or no call was issued yet. - """ - if self._state == self._STATE_NONE: - raise exc.ProgrammingError("No query yet") - - # Sleep until we're done or we have some data to return - self._fetch_while(lambda: not self._data and self._state != self._STATE_FINISHED) - - if not self._data: - return None - else: - self._rownumber += 1 - return self._data.popleft() - - def fetchmany(self, size=None): - """Fetch the next set of rows of a query result, returning a sequence of sequences (e.g. a - list of tuples). An empty sequence is returned when no more rows are available. - - The number of rows to fetch per call is specified by the parameter. If it is not given, the - cursor's arraysize determines the number of rows to be fetched. The method should try to - fetch as many rows as indicated by the size parameter. If this is not possible due to the - specified number of rows not being available, fewer rows may be returned. - - An :py:class:`~pyhive.exc.Error` (or subclass) exception is raised if the previous call to - :py:meth:`execute` did not produce any result set or no call was issued yet. - """ - if size is None: - size = self.arraysize - return list(islice(iter(self.fetchone, None), size)) - - def fetchall(self): - """Fetch all (remaining) rows of a query result, returning them as a sequence of sequences - (e.g. a list of tuples). - - An :py:class:`~pyhive.exc.Error` (or subclass) exception is raised if the previous call to - :py:meth:`execute` did not produce any result set or no call was issued yet. - """ - return list(iter(self.fetchone, None)) - - @property - def arraysize(self): - """This read/write attribute specifies the number of rows to fetch at a time with - :py:meth:`fetchmany`. It defaults to 1 meaning to fetch a single row at a time. - """ - return self._arraysize - - @arraysize.setter - def arraysize(self, value): - self._arraysize = value - - def setinputsizes(self, sizes): - """Does nothing by default""" - pass - - def setoutputsize(self, size, column=None): - """Does nothing by default""" - pass - - # - # Optional DB API Extensions - # - - @property - def rownumber(self): - """This read-only attribute should provide the current 0-based index of the cursor in the - result set. - - The index can be seen as index of the cursor in a sequence (the result set). The next fetch - operation will fetch the row indexed by ``rownumber`` in that sequence. - """ - return self._rownumber - - def __next__(self): - """Return the next row from the currently executing SQL statement using the same semantics - as :py:meth:`fetchone`. A ``StopIteration`` exception is raised when the result set is - exhausted. - """ - one = self.fetchone() - if one is None: - raise StopIteration - else: - return one - - next = __next__ - - def __iter__(self): - """Return self to make cursors compatible to the iteration protocol.""" - return self - - -class DBAPITypeObject(object): - # Taken from http://www.python.org/dev/peps/pep-0249/#implementation-hints - def __init__(self, *values): - self.values = values - - def __cmp__(self, other): - if other in self.values: - return 0 - if other < self.values: - return 1 - else: - return -1 - - -class ParamEscaper(object): - _DATE_FORMAT = "%Y-%m-%d" - _TIME_FORMAT = "%H:%M:%S.%f" - _DATETIME_FORMAT = "{} {}".format(_DATE_FORMAT, _TIME_FORMAT) - - def escape_args(self, parameters): - if isinstance(parameters, dict): - return {k: self.escape_item(v) for k, v in parameters.items()} - elif isinstance(parameters, (list, tuple)): - return tuple(self.escape_item(x) for x in parameters) - else: - raise exc.ProgrammingError("Unsupported param format: {}".format(parameters)) - - def escape_number(self, item): - return item - - def escape_string(self, item): - # Need to decode UTF-8 because of old sqlalchemy. - # Newer SQLAlchemy checks dialect.supports_unicode_binds before encoding Unicode strings - # as byte strings. The old version always encodes Unicode as byte strings, which breaks - # string formatting here. - if isinstance(item, bytes): - item = item.decode('utf-8') - # This is good enough when backslashes are literal, newlines are just followed, and the way - # to escape a single quote is to put two single quotes. - # (i.e. only special character is single quote) - return "'{}'".format(item.replace("'", "''")) - - def escape_sequence(self, item): - l = map(str, map(self.escape_item, item)) - return '(' + ','.join(l) + ')' - - def escape_datetime(self, item, format, cutoff=0): - dt_str = item.strftime(format) - formatted = dt_str[:-cutoff] if cutoff and format.endswith(".%f") else dt_str - return "'{}'".format(formatted) - - def escape_item(self, item): - if item is None: - return 'NULL' - elif isinstance(item, (int, float)): - return self.escape_number(item) - elif isinstance(item, basestring): - return self.escape_string(item) - elif isinstance(item, Iterable): - return self.escape_sequence(item) - elif isinstance(item, datetime.datetime): - return self.escape_datetime(item, self._DATETIME_FORMAT) - elif isinstance(item, datetime.date): - return self.escape_datetime(item, self._DATE_FORMAT) - else: - raise exc.ProgrammingError("Unsupported object {}".format(item)) - - -class UniversalSet(object): - """set containing everything""" - def __contains__(self, item): - return True diff --git a/src/chronify/_vendor/kyuubi/pyhive/exc.py b/src/chronify/_vendor/kyuubi/pyhive/exc.py deleted file mode 100644 index 931cf21..0000000 --- a/src/chronify/_vendor/kyuubi/pyhive/exc.py +++ /dev/null @@ -1,72 +0,0 @@ -""" -Package private common utilities. Do not use directly. -""" -from __future__ import absolute_import -from __future__ import unicode_literals - -__all__ = [ - 'Error', 'Warning', 'InterfaceError', 'DatabaseError', 'InternalError', 'OperationalError', - 'ProgrammingError', 'DataError', 'NotSupportedError', -] - - -class Error(Exception): - """Exception that is the base class of all other error exceptions. - - You can use this to catch all errors with one single except statement. - """ - pass - - -class Warning(Exception): - """Exception raised for important warnings like data truncations while inserting, etc.""" - pass - - -class InterfaceError(Error): - """Exception raised for errors that are related to the database interface rather than the - database itself. - """ - pass - - -class DatabaseError(Error): - """Exception raised for errors that are related to the database.""" - pass - - -class InternalError(DatabaseError): - """Exception raised when the database encounters an internal error, e.g. the cursor is not valid - anymore, the transaction is out of sync, etc.""" - pass - - -class OperationalError(DatabaseError): - """Exception raised for errors that are related to the database's operation and not necessarily - under the control of the programmer, e.g. an unexpected disconnect occurs, the data source name - is not found, a transaction could not be processed, a memory allocation error occurred during - processing, etc. - """ - pass - - -class ProgrammingError(DatabaseError): - """Exception raised for programming errors, e.g. table not found or already exists, syntax error - in the SQL statement, wrong number of parameters specified, etc. - """ - pass - - -class DataError(DatabaseError): - """Exception raised for errors that are due to problems with the processed data like division by - zero, numeric value out of range, etc. - """ - pass - - -class NotSupportedError(DatabaseError): - """Exception raised in case a method or database API was used which is not supported by the - database, e.g. requesting a ``.rollback()`` on a connection that does not support transaction or - has transactions turned off. - """ - pass diff --git a/src/chronify/_vendor/kyuubi/pyhive/hive.py b/src/chronify/_vendor/kyuubi/pyhive/hive.py deleted file mode 100644 index d6f8080..0000000 --- a/src/chronify/_vendor/kyuubi/pyhive/hive.py +++ /dev/null @@ -1,620 +0,0 @@ -"""DB-API implementation backed by HiveServer2 (Thrift API) - -See http://www.python.org/dev/peps/pep-0249/ - -Many docstrings in this file are based on the PEP, which is in the public domain. -""" - -from __future__ import absolute_import -from __future__ import unicode_literals - -import base64 -import datetime -import re -from decimal import Decimal -from ssl import CERT_NONE, CERT_OPTIONAL, CERT_REQUIRED, create_default_context - - -from TCLIService import TCLIService -from TCLIService import constants -from TCLIService import ttypes -from pyhive import common -from pyhive.common import DBAPITypeObject -# Make all exceptions visible in this module per DB-API -from pyhive.exc import * # noqa -from builtins import range -import contextlib -from future.utils import iteritems -import getpass -import logging -import sys -import thrift.transport.THttpClient -import thrift.protocol.TBinaryProtocol -import thrift.transport.TSocket -import thrift.transport.TTransport - -# PEP 249 module globals -apilevel = '2.0' -threadsafety = 2 # Threads may share the module and connections. -paramstyle = 'pyformat' # Python extended format codes, e.g. ...WHERE name=%(name)s - -_logger = logging.getLogger(__name__) - -_TIMESTAMP_PATTERN = re.compile(r'(\d+-\d+-\d+ \d+:\d+:\d+(\.\d{,6})?)') - -ssl_cert_parameter_map = { - "none": CERT_NONE, - "optional": CERT_OPTIONAL, - "required": CERT_REQUIRED, -} - - -def get_sasl_client(host, sasl_auth, service=None, username=None, password=None): - import sasl - sasl_client = sasl.Client() - sasl_client.setAttr('host', host) - - if sasl_auth == 'GSSAPI': - sasl_client.setAttr('service', service) - elif sasl_auth == 'PLAIN': - sasl_client.setAttr('username', username) - sasl_client.setAttr('password', password) - else: - raise ValueError("sasl_auth only supports GSSAPI and PLAIN") - - sasl_client.init() - return sasl_client - - -def get_pure_sasl_client(host, sasl_auth, service=None, username=None, password=None): - from pyhive.sasl_compat import PureSASLClient - - if sasl_auth == 'GSSAPI': - sasl_kwargs = {'service': service} - elif sasl_auth == 'PLAIN': - sasl_kwargs = {'username': username, 'password': password} - else: - raise ValueError("sasl_auth only supports GSSAPI and PLAIN") - - return PureSASLClient(host=host, **sasl_kwargs) - - -def get_installed_sasl(host, sasl_auth, service=None, username=None, password=None): - try: - return get_sasl_client(host=host, sasl_auth=sasl_auth, service=service, username=username, password=password) - # The sasl library is available - except ImportError: - # Fallback to pure-sasl library - return get_pure_sasl_client(host=host, sasl_auth=sasl_auth, service=service, username=username, password=password) - - -def _parse_timestamp(value): - if value: - match = _TIMESTAMP_PATTERN.match(value) - if match: - if match.group(2): - format = '%Y-%m-%d %H:%M:%S.%f' - # use the pattern to truncate the value - value = match.group() - else: - format = '%Y-%m-%d %H:%M:%S' - value = datetime.datetime.strptime(value, format) - else: - raise Exception( - 'Cannot convert "{}" into a datetime'.format(value)) - else: - value = None - return value - - -TYPES_CONVERTER = {"DECIMAL_TYPE": Decimal, - "TIMESTAMP_TYPE": _parse_timestamp} - - -class HiveParamEscaper(common.ParamEscaper): - def escape_string(self, item): - # backslashes and single quotes need to be escaped - # TODO verify against parser - # Need to decode UTF-8 because of old sqlalchemy. - # Newer SQLAlchemy checks dialect.supports_unicode_binds before encoding Unicode strings - # as byte strings. The old version always encodes Unicode as byte strings, which breaks - # string formatting here. - if isinstance(item, bytes): - item = item.decode('utf-8') - return "'{}'".format( - item - .replace('\\', '\\\\') - .replace("'", "\\'") - .replace('\r', '\\r') - .replace('\n', '\\n') - .replace('\t', '\\t') - ) - - -_escaper = HiveParamEscaper() - - -def connect(*args, **kwargs): - """Constructor for creating a connection to the database. See class :py:class:`Connection` for - arguments. - - :returns: a :py:class:`Connection` object. - """ - return Connection(*args, **kwargs) - - -class Connection(object): - """Wraps a Thrift session""" - - def __init__( - self, - host=None, - port=None, - scheme=None, - username=None, - database='default', - auth=None, - configuration=None, - kerberos_service_name=None, - password=None, - check_hostname=None, - ssl_cert=None, - thrift_transport=None, - ssl_context=None - ): - """Connect to HiveServer2 - - :param host: What host HiveServer2 runs on - :param port: What port HiveServer2 runs on. Defaults to 10000. - :param auth: The value of hive.server2.authentication used by HiveServer2. - Defaults to ``NONE``. - :param configuration: A dictionary of Hive settings (functionally same as the `set` command) - :param kerberos_service_name: Use with auth='KERBEROS' only - :param password: Use with auth='LDAP' or auth='CUSTOM' only - :param thrift_transport: A ``TTransportBase`` for custom advanced usage. - Incompatible with host, port, auth, kerberos_service_name, and password. - :param ssl_context: A custom SSL context to use for HTTPS connections. If provided, - this overrides check_hostname and ssl_cert parameters. - The way to support LDAP and GSSAPI is originated from cloudera/Impyla: - https://github.com/cloudera/impyla/blob/255b07ed973d47a3395214ed92d35ec0615ebf62 - /impala/_thrift_api.py#L152-L160 - """ - if scheme in ("https", "http") and thrift_transport is None: - port = port or 1000 - if scheme == "https": - if ssl_context is None: - ssl_context = create_default_context() - ssl_context.check_hostname = check_hostname == "true" - ssl_cert = ssl_cert or "none" - ssl_context.verify_mode = ssl_cert_parameter_map.get(ssl_cert, CERT_NONE) - thrift_transport = thrift.transport.THttpClient.THttpClient( - uri_or_host="{scheme}://{host}:{port}/cliservice/".format( - scheme=scheme, host=host, port=port - ), - ssl_context=ssl_context, - ) - - if auth in ("BASIC", "NOSASL", "NONE", None): - # Always needs the Authorization header - self._set_authorization_header(thrift_transport, username, password) - elif auth == "KERBEROS" and kerberos_service_name: - self._set_kerberos_header(thrift_transport, kerberos_service_name, host) - else: - raise ValueError( - "Authentication is not valid use one of:" - "BASIC, NOSASL, KERBEROS, NONE" - ) - host, port, auth, kerberos_service_name, password = ( - None, None, None, None, None - ) - - username = username or getpass.getuser() - configuration = configuration or {} - - if (password is not None) != (auth in ('LDAP', 'CUSTOM')): - raise ValueError("Password should be set if and only if in LDAP or CUSTOM mode; " - "Remove password or use one of those modes") - if (kerberos_service_name is not None) != (auth == 'KERBEROS'): - raise ValueError("kerberos_service_name should be set if and only if in KERBEROS mode") - if thrift_transport is not None: - has_incompatible_arg = ( - host is not None - or port is not None - or auth is not None - or kerberos_service_name is not None - or password is not None - ) - if has_incompatible_arg: - raise ValueError("thrift_transport cannot be used with " - "host/port/auth/kerberos_service_name/password") - - if thrift_transport is not None: - self._transport = thrift_transport - else: - if port is None: - port = 10000 - if auth is None: - auth = 'NONE' - socket = thrift.transport.TSocket.TSocket(host, port) - if auth == 'NOSASL': - # NOSASL corresponds to hive.server2.authentication=NOSASL in hive-site.xml - self._transport = thrift.transport.TTransport.TBufferedTransport(socket) - elif auth in ('LDAP', 'KERBEROS', 'NONE', 'CUSTOM'): - # Defer import so package dependency is optional - import thrift_sasl - - if auth == 'KERBEROS': - # KERBEROS mode in hive.server2.authentication is GSSAPI in sasl library - sasl_auth = 'GSSAPI' - else: - sasl_auth = 'PLAIN' - if password is None: - # Password doesn't matter in NONE mode, just needs to be nonempty. - password = 'x' - - self._transport = thrift_sasl.TSaslClientTransport(lambda: get_installed_sasl(host=host, sasl_auth=sasl_auth, service=kerberos_service_name, username=username, password=password), sasl_auth, socket) - else: - # All HS2 config options: - # https://cwiki.apache.org/confluence/display/Hive/Setting+Up+HiveServer2#SettingUpHiveServer2-Configuration - # PAM currently left to end user via thrift_transport option. - raise NotImplementedError( - "Only NONE, NOSASL, LDAP, KERBEROS, CUSTOM " - "authentication are supported, got {}".format(auth)) - - protocol = thrift.protocol.TBinaryProtocol.TBinaryProtocol(self._transport) - self._client = TCLIService.Client(protocol) - # oldest version that still contains features we care about - # "V6 uses binary type for binary payload (was string) and uses columnar result set" - protocol_version = ttypes.TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V6 - - try: - self._transport.open() - open_session_req = ttypes.TOpenSessionReq( - client_protocol=protocol_version, - configuration=configuration, - username=username, - ) - response = self._client.OpenSession(open_session_req) - _check_status(response) - assert response.sessionHandle is not None, "Expected a session from OpenSession" - self._sessionHandle = response.sessionHandle - assert response.serverProtocolVersion == protocol_version, \ - "Unable to handle protocol version {}".format(response.serverProtocolVersion) - with contextlib.closing(self.cursor()) as cursor: - cursor.execute('USE `{}`'.format(database)) - except: - self._transport.close() - raise - - @staticmethod - def _set_authorization_header(transport, username=None, password=None): - username = username or "user" - password = password or "pass" - auth_credentials = "{username}:{password}".format( - username=username, password=password - ).encode("UTF-8") - auth_credentials_base64 = base64.standard_b64encode(auth_credentials).decode( - "UTF-8" - ) - transport.setCustomHeaders( - { - "Authorization": "Basic {auth_credentials_base64}".format( - auth_credentials_base64=auth_credentials_base64 - ) - } - ) - - @staticmethod - def _set_kerberos_header(transport, kerberos_service_name, host): - import kerberos - - __, krb_context = kerberos.authGSSClientInit( - service="{kerberos_service_name}@{host}".format( - kerberos_service_name=kerberos_service_name, host=host - ) - ) - kerberos.authGSSClientClean(krb_context, "") - kerberos.authGSSClientStep(krb_context, "") - auth_header = kerberos.authGSSClientResponse(krb_context) - - transport.setCustomHeaders( - { - "Authorization": "Negotiate {auth_header}".format( - auth_header=auth_header - ) - } - ) - - def __enter__(self): - """Transport should already be opened by __init__""" - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - """Call close""" - self.close() - - def close(self): - """Close the underlying session and Thrift transport""" - req = ttypes.TCloseSessionReq(sessionHandle=self._sessionHandle) - response = self._client.CloseSession(req) - self._transport.close() - _check_status(response) - - def commit(self): - """Hive does not support transactions, so this does nothing.""" - pass - - def cursor(self, *args, **kwargs): - """Return a new :py:class:`Cursor` object using the connection.""" - return Cursor(self, *args, **kwargs) - - @property - def client(self): - return self._client - - @property - def sessionHandle(self): - return self._sessionHandle - - def rollback(self): - raise NotSupportedError("Hive does not have transactions") # pragma: no cover - - -class Cursor(common.DBAPICursor): - """These objects represent a database cursor, which is used to manage the context of a fetch - operation. - - Cursors are not isolated, i.e., any changes done to the database by a cursor are immediately - visible by other cursors or connections. - """ - - def __init__(self, connection, arraysize=1000): - self._operationHandle = None - super(Cursor, self).__init__() - self._arraysize = arraysize - self._connection = connection - - def _reset_state(self): - """Reset state about the previous query in preparation for running another query""" - super(Cursor, self)._reset_state() - self._description = None - if self._operationHandle is not None: - request = ttypes.TCloseOperationReq(self._operationHandle) - try: - response = self._connection.client.CloseOperation(request) - _check_status(response) - finally: - self._operationHandle = None - - @property - def arraysize(self): - return self._arraysize - - @arraysize.setter - def arraysize(self, value): - """Array size cannot be None, and should be an integer""" - default_arraysize = 1000 - try: - self._arraysize = int(value) or default_arraysize - except TypeError: - self._arraysize = default_arraysize - - @property - def description(self): - """This read-only attribute is a sequence of 7-item sequences. - - Each of these sequences contains information describing one result column: - - - name - - type_code - - display_size (None in current implementation) - - internal_size (None in current implementation) - - precision (None in current implementation) - - scale (None in current implementation) - - null_ok (always True in current implementation) - - This attribute will be ``None`` for operations that do not return rows or if the cursor has - not had an operation invoked via the :py:meth:`execute` method yet. - - The ``type_code`` can be interpreted by comparing it to the Type Objects specified in the - section below. - """ - if self._operationHandle is None or not self._operationHandle.hasResultSet: - return None - if self._description is None: - req = ttypes.TGetResultSetMetadataReq(self._operationHandle) - response = self._connection.client.GetResultSetMetadata(req) - _check_status(response) - columns = response.schema.columns - self._description = [] - for col in columns: - primary_type_entry = col.typeDesc.types[0] - if primary_type_entry.primitiveEntry is None: - # All fancy stuff maps to string - type_code = ttypes.TTypeId._VALUES_TO_NAMES[ttypes.TTypeId.STRING_TYPE] - else: - type_id = primary_type_entry.primitiveEntry.type - try: - type_code = ttypes.TTypeId._VALUES_TO_NAMES[type_id] - except KeyError: - type_code = None - self._description.append(( - col.columnName.decode('utf-8') if sys.version_info[0] == 2 else col.columnName, - type_code.decode('utf-8') if sys.version_info[0] == 2 else type_code, - None, None, None, None, True - )) - return self._description - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - self.close() - - def close(self): - """Close the operation handle""" - self._reset_state() - - def execute(self, operation, parameters=None, **kwargs): - """Prepare and execute a database operation (query or command). - - Return values are not defined. - """ - # backward compatibility with Python < 3.7 - for kw in ['async', 'async_']: - if kw in kwargs: - async_ = kwargs[kw] - break - else: - async_ = False - - # Prepare statement - if parameters is None: - sql = operation - else: - sql = operation % _escaper.escape_args(parameters) - - self._reset_state() - - self._state = self._STATE_RUNNING - _logger.info('%s', sql) - - req = ttypes.TExecuteStatementReq(self._connection.sessionHandle, - sql, runAsync=async_) - _logger.debug(req) - response = self._connection.client.ExecuteStatement(req) - _check_status(response) - self._operationHandle = response.operationHandle - - def cancel(self): - req = ttypes.TCancelOperationReq( - operationHandle=self._operationHandle, - ) - response = self._connection.client.CancelOperation(req) - _check_status(response) - - def _fetch_more(self): - """Send another TFetchResultsReq and update state""" - assert(self._state == self._STATE_RUNNING), "Should be running when in _fetch_more" - assert(self._operationHandle is not None), "Should have an op handle in _fetch_more" - if not self._operationHandle.hasResultSet: - raise ProgrammingError("No result set") - req = ttypes.TFetchResultsReq( - operationHandle=self._operationHandle, - orientation=ttypes.TFetchOrientation.FETCH_NEXT, - maxRows=self.arraysize, - ) - response = self._connection.client.FetchResults(req) - _check_status(response) - schema = self.description - assert not response.results.rows, 'expected data in columnar format' - has_new_data = False - if response.results.columns: - columns = [_unwrap_column(col, col_schema[1]) for col, col_schema in - zip(response.results.columns, schema)] - new_data = list(zip(*columns)) - self._data += new_data - has_new_data = (True if new_data else False) - # response.hasMoreRows seems to always be False, so we instead check the number of rows - # https://github.com/apache/hive/blob/release-1.2.1/service/src/java/org/apache/hive/service/cli/thrift/ThriftCLIService.java#L678 - # if not response.hasMoreRows: - if not has_new_data: - self._state = self._STATE_FINISHED - - def poll(self, get_progress_update=True): - """Poll for and return the raw status data provided by the Hive Thrift REST API. - :returns: ``ttypes.TGetOperationStatusResp`` - :raises: ``ProgrammingError`` when no query has been started - .. note:: - This is not a part of DB-API. - """ - if self._state == self._STATE_NONE: - raise ProgrammingError("No query yet") - - req = ttypes.TGetOperationStatusReq( - operationHandle=self._operationHandle, - getProgressUpdate=get_progress_update, - ) - response = self._connection.client.GetOperationStatus(req) - _check_status(response) - - return response - - def fetch_logs(self): - """Retrieve the logs produced by the execution of the query. - Can be called multiple times to fetch the logs produced after the previous call. - :returns: list - :raises: ``ProgrammingError`` when no query has been started - .. note:: - This is not a part of DB-API. - """ - if self._state == self._STATE_NONE: - raise ProgrammingError("No query yet") - - try: # Older Hive instances require logs to be retrieved using GetLog - req = ttypes.TGetLogReq(operationHandle=self._operationHandle) - logs = self._connection.client.GetLog(req).log.splitlines() - except ttypes.TApplicationException as e: # Otherwise, retrieve logs using newer method - if e.type != ttypes.TApplicationException.UNKNOWN_METHOD: - raise - logs = [] - while True: - req = ttypes.TFetchResultsReq( - operationHandle=self._operationHandle, - orientation=ttypes.TFetchOrientation.FETCH_NEXT, - maxRows=self.arraysize, - fetchType=1 # 0: results, 1: logs - ) - response = self._connection.client.FetchResults(req) - _check_status(response) - assert not response.results.rows, 'expected data in columnar format' - new_logs = '' - if response.results.columns: - new_logs = _unwrap_column(response.results.columns[0]) - logs += new_logs - - if not new_logs: - break - - return logs - - -# -# Type Objects and Constructors -# - - -for type_id in constants.PRIMITIVE_TYPES: - name = ttypes.TTypeId._VALUES_TO_NAMES[type_id] - setattr(sys.modules[__name__], name, DBAPITypeObject([name])) - - -# -# Private utilities -# - - -def _unwrap_column(col, type_=None): - """Return a list of raw values from a TColumn instance.""" - for attr, wrapper in iteritems(col.__dict__): - if wrapper is not None: - result = wrapper.values - nulls = wrapper.nulls # bit set describing what's null - assert isinstance(nulls, bytes) - for i, char in enumerate(nulls): - byte = ord(char) if sys.version_info[0] == 2 else char - for b in range(8): - if byte & (1 << b): - result[i * 8 + b] = None - converter = TYPES_CONVERTER.get(type_, None) - if converter and type_: - result = [converter(row) if row else row for row in result] - return result - raise DataError("Got empty column value {}".format(col)) # pragma: no cover - - -def _check_status(response): - """Raise an OperationalError if the status is not success""" - _logger.debug(response) - if response.status.statusCode != ttypes.TStatusCode.SUCCESS_STATUS: - raise OperationalError(response) diff --git a/src/chronify/_vendor/kyuubi/pyhive/presto.py b/src/chronify/_vendor/kyuubi/pyhive/presto.py deleted file mode 100644 index 3217f4c..0000000 --- a/src/chronify/_vendor/kyuubi/pyhive/presto.py +++ /dev/null @@ -1,367 +0,0 @@ -"""DB-API implementation backed by Presto - -See http://www.python.org/dev/peps/pep-0249/ - -Many docstrings in this file are based on the PEP, which is in the public domain. -""" - -from __future__ import absolute_import -from __future__ import unicode_literals - -from builtins import object -from decimal import Decimal - -from pyhive import common -from pyhive.common import DBAPITypeObject -# Make all exceptions visible in this module per DB-API -from pyhive.exc import * # noqa -import base64 -import getpass -import datetime -import logging -import requests -from requests.auth import HTTPBasicAuth -import os - -try: # Python 3 - import urllib.parse as urlparse -except ImportError: # Python 2 - import urlparse - - -# PEP 249 module globals -apilevel = '2.0' -threadsafety = 2 # Threads may share the module and connections. -paramstyle = 'pyformat' # Python extended format codes, e.g. ...WHERE name=%(name)s - -_logger = logging.getLogger(__name__) - -TYPES_CONVERTER = { - "decimal": Decimal, - # As of Presto 0.69, binary data is returned as the varbinary type in base64 format - "varbinary": base64.b64decode -} - -class PrestoParamEscaper(common.ParamEscaper): - def escape_datetime(self, item, format): - _type = "timestamp" if isinstance(item, datetime.datetime) else "date" - formatted = super(PrestoParamEscaper, self).escape_datetime(item, format, 3) - return "{} {}".format(_type, formatted) - - -_escaper = PrestoParamEscaper() - - -def connect(*args, **kwargs): - """Constructor for creating a connection to the database. See class :py:class:`Connection` for - arguments. - - :returns: a :py:class:`Connection` object. - """ - return Connection(*args, **kwargs) - - -class Connection(object): - """Presto does not have a notion of a persistent connection. - - Thus, these objects are small stateless factories for cursors, which do all the real work. - """ - - def __init__(self, *args, **kwargs): - self._args = args - self._kwargs = kwargs - - def close(self): - """Presto does not have anything to close""" - # TODO cancel outstanding queries? - pass - - def commit(self): - """Presto does not support transactions""" - pass - - def cursor(self): - """Return a new :py:class:`Cursor` object using the connection.""" - return Cursor(*self._args, **self._kwargs) - - def rollback(self): - raise NotSupportedError("Presto does not have transactions") # pragma: no cover - - -class Cursor(common.DBAPICursor): - """These objects represent a database cursor, which is used to manage the context of a fetch - operation. - - Cursors are not isolated, i.e., any changes done to the database by a cursor are immediately - visible by other cursors or connections. - """ - - def __init__(self, host, port='8080', username=None, principal_username=None, catalog='hive', - schema='default', poll_interval=1, source='pyhive', session_props=None, - protocol='http', password=None, requests_session=None, requests_kwargs=None, - KerberosRemoteServiceName=None, KerberosPrincipal=None, - KerberosConfigPath=None, KerberosKeytabPath=None, - KerberosCredentialCachePath=None, KerberosUseCanonicalHostname=None): - """ - :param host: hostname to connect to, e.g. ``presto.example.com`` - :param port: int -- port, defaults to 8080 - :param username: string -- defaults to system user name - :param principal_username: string -- defaults to ``username`` argument if it exists, - else defaults to system user name - :param catalog: string -- defaults to ``hive`` - :param schema: string -- defaults to ``default`` - :param poll_interval: float -- how often to ask the Presto REST interface for a progress - update, defaults to a second - :param source: string -- arbitrary identifier (shows up in the Presto monitoring page) - :param protocol: string -- network protocol, valid options are ``http`` and ``https``. - defaults to ``http`` - :param password: string -- Deprecated. Defaults to ``None``. - Using BasicAuth, requires ``https``. - Prefer ``requests_kwargs={'auth': HTTPBasicAuth(username, password)}``. - May not be specified with ``requests_kwargs['auth']``. - :param requests_session: a ``requests.Session`` object for advanced usage. If absent, this - class will use the default requests behavior of making a new session per HTTP request. - Caller is responsible for closing session. - :param requests_kwargs: Additional ``**kwargs`` to pass to requests - :param KerberosRemoteServiceName: string -- Presto coordinator Kerberos service name. - This parameter is required for Kerberos authentiation. - :param KerberosPrincipal: string -- The principal to use when authenticating to - the Presto coordinator. - :param KerberosConfigPath: string -- Kerberos configuration file. - (default: /etc/krb5.conf) - :param KerberosKeytabPath: string -- Kerberos keytab file. - :param KerberosCredentialCachePath: string -- Kerberos credential cache. - :param KerberosUseCanonicalHostname: boolean -- Use the canonical hostname of the - Presto coordinator for the Kerberos service principal by first resolving the - hostname to an IP address and then doing a reverse DNS lookup for that IP address. - This is enabled by default. - """ - super(Cursor, self).__init__(poll_interval) - # Config - self._host = host - self._port = port - """ - Presto User Impersonation: https://docs.starburstdata.com/latest/security/impersonation.html - - User impersonation allows the execution of queries in Presto based on principal_username - argument, instead of executing the query as the account which authenticated against Presto. - (Usually a service account) - - Allows for a service account to authenticate with Presto, and then leverage the - principal_username as the user Presto will execute the query as. This is required by - applications that leverage authentication methods like SAML, where the application has a - username, but not a password to still leverage user specific Presto Resource Groups and - Authorization rules that would not be applied when only using a shared service account. - This also allows auditing of who is executing a query in these environments, instead of - having all queryes run by the shared service account. - """ - self._username = principal_username or username or getpass.getuser() - self._catalog = catalog - self._schema = schema - self._arraysize = 1 - self._poll_interval = poll_interval - self._source = source - self._session_props = session_props if session_props is not None else {} - self.last_query_id = None - - if protocol not in ('http', 'https'): - raise ValueError("Protocol must be http/https, was {!r}".format(protocol)) - self._protocol = protocol - - self._requests_session = requests_session or requests - - requests_kwargs = dict(requests_kwargs) if requests_kwargs is not None else {} - - if KerberosRemoteServiceName is not None: - from requests_kerberos import HTTPKerberosAuth, OPTIONAL - - hostname_override = None - if KerberosUseCanonicalHostname is not None \ - and KerberosUseCanonicalHostname.lower() == 'false': - hostname_override = host - if KerberosConfigPath is not None: - os.environ['KRB5_CONFIG'] = KerberosConfigPath - if KerberosKeytabPath is not None: - os.environ['KRB5_CLIENT_KTNAME'] = KerberosKeytabPath - if KerberosCredentialCachePath is not None: - os.environ['KRB5CCNAME'] = KerberosCredentialCachePath - - requests_kwargs['auth'] = HTTPKerberosAuth(mutual_authentication=OPTIONAL, - principal=KerberosPrincipal, - service=KerberosRemoteServiceName, - hostname_override=hostname_override) - - else: - if password is not None and 'auth' in requests_kwargs: - raise ValueError("Cannot use both password and requests_kwargs authentication") - for k in ('method', 'url', 'data', 'headers'): - if k in requests_kwargs: - raise ValueError("Cannot override requests argument {}".format(k)) - if password is not None: - requests_kwargs['auth'] = HTTPBasicAuth(username, password) - if protocol != 'https': - raise ValueError("Protocol must be https when passing a password") - self._requests_kwargs = requests_kwargs - - self._reset_state() - - def _reset_state(self): - """Reset state about the previous query in preparation for running another query""" - super(Cursor, self)._reset_state() - self._nextUri = None - self._columns = None - - @property - def description(self): - """This read-only attribute is a sequence of 7-item sequences. - - Each of these sequences contains information describing one result column: - - - name - - type_code - - display_size (None in current implementation) - - internal_size (None in current implementation) - - precision (None in current implementation) - - scale (None in current implementation) - - null_ok (always True in current implementation) - - The ``type_code`` can be interpreted by comparing it to the Type Objects specified in the - section below. - """ - # Sleep until we're done or we got the columns - self._fetch_while( - lambda: self._columns is None and - self._state not in (self._STATE_NONE, self._STATE_FINISHED) - ) - if self._columns is None: - return None - return [ - # name, type_code, display_size, internal_size, precision, scale, null_ok - (col['name'], col['type'], None, None, None, None, True) - for col in self._columns - ] - - def execute(self, operation, parameters=None): - """Prepare and execute a database operation (query or command). - - Return values are not defined. - """ - headers = { - 'X-Presto-Catalog': self._catalog, - 'X-Presto-Schema': self._schema, - 'X-Presto-Source': self._source, - 'X-Presto-User': self._username, - } - - if self._session_props: - headers['X-Presto-Session'] = ','.join( - '{}={}'.format(propname, propval) - for propname, propval in self._session_props.items() - ) - - # Prepare statement - if parameters is None: - sql = operation - else: - sql = operation % _escaper.escape_args(parameters) - - self._reset_state() - - self._state = self._STATE_RUNNING - url = urlparse.urlunparse(( - self._protocol, - '{}:{}'.format(self._host, self._port), '/v1/statement', None, None, None)) - _logger.info('%s', sql) - _logger.debug("Headers: %s", headers) - response = self._requests_session.post( - url, data=sql.encode('utf-8'), headers=headers, **self._requests_kwargs) - self._process_response(response) - - def cancel(self): - if self._state == self._STATE_NONE: - raise ProgrammingError("No query yet") - if self._nextUri is None: - assert self._state == self._STATE_FINISHED, "Should be finished if nextUri is None" - return - - response = self._requests_session.delete(self._nextUri, **self._requests_kwargs) - if response.status_code != requests.codes.no_content: - fmt = "Unexpected status code after cancel {}\n{}" - raise OperationalError(fmt.format(response.status_code, response.content)) - - self._state = self._STATE_FINISHED - self._nextUri = None - - def poll(self): - """Poll for and return the raw status data provided by the Presto REST API. - - :returns: dict -- JSON status information or ``None`` if the query is done - :raises: ``ProgrammingError`` when no query has been started - - .. note:: - This is not a part of DB-API. - """ - if self._state == self._STATE_NONE: - raise ProgrammingError("No query yet") - if self._nextUri is None: - assert self._state == self._STATE_FINISHED, "Should be finished if nextUri is None" - return None - response = self._requests_session.get(self._nextUri, **self._requests_kwargs) - self._process_response(response) - return response.json() - - def _fetch_more(self): - """Fetch the next URI and update state""" - self._process_response(self._requests_session.get(self._nextUri, **self._requests_kwargs)) - - def _process_data(self, rows): - for i, col in enumerate(self.description): - col_type = col[1].split("(")[0].lower() - if col_type in TYPES_CONVERTER: - for row in rows: - if row[i] is not None: - row[i] = TYPES_CONVERTER[col_type](row[i]) - - def _process_response(self, response): - """Given the JSON response from Presto's REST API, update the internal state with the next - URI and any data from the response - """ - # TODO handle HTTP 503 - if response.status_code != requests.codes.ok: - fmt = "Unexpected status code {}\n{}" - raise OperationalError(fmt.format(response.status_code, response.content)) - - response_json = response.json() - _logger.debug("Got response %s", response_json) - assert self._state == self._STATE_RUNNING, "Should be running if processing response" - self._nextUri = response_json.get('nextUri') - self._columns = response_json.get('columns') - if 'id' in response_json: - self.last_query_id = response_json['id'] - if 'X-Presto-Clear-Session' in response.headers: - propname = response.headers['X-Presto-Clear-Session'] - self._session_props.pop(propname, None) - if 'X-Presto-Set-Session' in response.headers: - propname, propval = response.headers['X-Presto-Set-Session'].split('=', 1) - self._session_props[propname] = propval - if 'data' in response_json: - assert self._columns - new_data = response_json['data'] - self._process_data(new_data) - self._data += map(tuple, new_data) - if 'nextUri' not in response_json: - self._state = self._STATE_FINISHED - if 'error' in response_json: - raise DatabaseError(response_json['error']) - - -# -# Type Objects and Constructors -# - - -# See types in presto-main/src/main/java/com/facebook/presto/tuple/TupleInfo.java -FIXED_INT_64 = DBAPITypeObject(['bigint']) -VARIABLE_BINARY = DBAPITypeObject(['varchar']) -DOUBLE = DBAPITypeObject(['double']) -BOOLEAN = DBAPITypeObject(['boolean']) diff --git a/src/chronify/_vendor/kyuubi/pyhive/sasl_compat.py b/src/chronify/_vendor/kyuubi/pyhive/sasl_compat.py deleted file mode 100644 index 19af6d2..0000000 --- a/src/chronify/_vendor/kyuubi/pyhive/sasl_compat.py +++ /dev/null @@ -1,56 +0,0 @@ -# Original source of this file is https://github.com/cloudera/impyla/blob/master/impala/sasl_compat.py -# which uses Apache-2.0 license as of 21 May 2023. -# This code was added to Impyla in 2016 as a compatibility layer to allow use of either python-sasl or pure-sasl -# via PR https://github.com/cloudera/impyla/pull/179 -# Even though thrift_sasl lists pure-sasl as dependency here https://github.com/cloudera/thrift_sasl/blob/master/setup.py#L34 -# but it still calls functions native to python-sasl in this file https://github.com/cloudera/thrift_sasl/blob/master/thrift_sasl/__init__.py#L82 -# Hence this code is required for the fallback to work. - - -from puresasl.client import SASLClient, SASLError -from contextlib import contextmanager - -@contextmanager -def error_catcher(self, Exc = Exception): - try: - self.error = None - yield - except Exc as e: - self.error = str(e) - - -class PureSASLClient(SASLClient): - def __init__(self, *args, **kwargs): - self.error = None - super(PureSASLClient, self).__init__(*args, **kwargs) - - def start(self, mechanism): - with error_catcher(self, SASLError): - if isinstance(mechanism, list): - self.choose_mechanism(mechanism) - else: - self.choose_mechanism([mechanism]) - return True, self.mechanism, self.process() - # else - return False, mechanism, None - - def encode(self, incoming): - with error_catcher(self): - return True, self.unwrap(incoming) - # else - return False, None - - def decode(self, outgoing): - with error_catcher(self): - return True, self.wrap(outgoing) - # else - return False, None - - def step(self, challenge=None): - with error_catcher(self): - return True, self.process(challenge) - # else - return False, None - - def getError(self): - return self.error diff --git a/src/chronify/_vendor/kyuubi/pyhive/sqlalchemy_hive.py b/src/chronify/_vendor/kyuubi/pyhive/sqlalchemy_hive.py deleted file mode 100644 index 66f7d75..0000000 --- a/src/chronify/_vendor/kyuubi/pyhive/sqlalchemy_hive.py +++ /dev/null @@ -1,435 +0,0 @@ -"""Integration between SQLAlchemy and Hive. - -Some code based on -https://github.com/zzzeek/sqlalchemy/blob/rel_0_5/lib/sqlalchemy/databases/sqlite.py -which is released under the MIT license. -""" - -from __future__ import absolute_import -from __future__ import unicode_literals - -import datetime -import decimal -import logging - -import re -from sqlalchemy import exc -from sqlalchemy.sql import text -try: - from sqlalchemy import processors -except ImportError: - # Required for SQLAlchemy>=2.0 - from sqlalchemy.engine import processors -from sqlalchemy import types -from sqlalchemy import util -# TODO shouldn't use mysql type -try: - from sqlalchemy.databases import mysql - mysql_tinyinteger = mysql.MSTinyInteger -except ImportError: - # Required for SQLAlchemy>2.0 - from sqlalchemy.dialects import mysql - mysql_tinyinteger = mysql.base.MSTinyInteger -from sqlalchemy.engine import default -from sqlalchemy.sql import compiler -from sqlalchemy.sql.compiler import SQLCompiler - -from pyhive import hive -from pyhive.common import UniversalSet - -from dateutil.parser import parse -from decimal import Decimal - -_logger = logging.getLogger(__name__) - -class HiveStringTypeBase(types.TypeDecorator): - """Translates strings returned by Thrift into something else""" - impl = types.String - - def process_bind_param(self, value, dialect): - raise NotImplementedError("Writing to Hive not supported") - - -class HiveDate(HiveStringTypeBase): - """Translates date strings to date objects""" - impl = types.DATE - - def process_result_value(self, value, dialect): - return processors.str_to_date(value) - - def result_processor(self, dialect, coltype): - def process(value): - if isinstance(value, datetime.datetime): - return value.date() - elif isinstance(value, datetime.date): - return value - elif value is not None: - return parse(value).date() - else: - return None - - return process - - def adapt(self, impltype, **kwargs): - return self.impl - - -class HiveTimestamp(HiveStringTypeBase): - """Translates timestamp strings to datetime objects""" - impl = types.TIMESTAMP - - def process_result_value(self, value, dialect): - return processors.str_to_datetime(value) - - def result_processor(self, dialect, coltype): - def process(value): - if isinstance(value, datetime.datetime): - return value - elif value is not None: - return parse(value) - else: - return None - - return process - - def adapt(self, impltype, **kwargs): - return self.impl - - -class HiveDecimal(HiveStringTypeBase): - """Translates strings to decimals""" - impl = types.DECIMAL - - def process_result_value(self, value, dialect): - if value is not None: - return decimal.Decimal(value) - else: - return None - - def result_processor(self, dialect, coltype): - def process(value): - if isinstance(value, Decimal): - return value - elif value is not None: - return Decimal(value) - else: - return None - - return process - - def adapt(self, impltype, **kwargs): - return self.impl - - -class HiveIdentifierPreparer(compiler.IdentifierPreparer): - # Just quote everything to make things simpler / easier to upgrade - reserved_words = UniversalSet() - - def __init__(self, dialect): - super(HiveIdentifierPreparer, self).__init__( - dialect, - initial_quote='`', - ) - - -_type_map = { - 'boolean': types.Boolean, - 'tinyint': mysql_tinyinteger, - 'smallint': types.SmallInteger, - 'int': types.Integer, - 'bigint': types.BigInteger, - 'float': types.Float, - 'double': types.Float, - 'string': types.String, - 'varchar': types.String, - 'char': types.String, - 'date': HiveDate, - 'timestamp': HiveTimestamp, - 'binary': types.String, - 'array': types.String, - 'map': types.String, - 'struct': types.String, - 'uniontype': types.String, - 'decimal': HiveDecimal, -} - - -class HiveCompiler(SQLCompiler): - def visit_concat_op_binary(self, binary, operator, **kw): - return "concat(%s, %s)" % (self.process(binary.left), self.process(binary.right)) - - def visit_insert(self, *args, **kwargs): - result = super(HiveCompiler, self).visit_insert(*args, **kwargs) - # Massage the result into Hive's format - # INSERT INTO `pyhive_test_database`.`test_table` (`a`) SELECT ... - # => - # INSERT INTO TABLE `pyhive_test_database`.`test_table` SELECT ... - regex = r'^(INSERT INTO) ([^\s]+) \([^\)]*\)' - assert re.search(regex, result), "Unexpected visit_insert result: {}".format(result) - return re.sub(regex, r'\1 TABLE \2', result) - - def visit_column(self, *args, **kwargs): - result = super(HiveCompiler, self).visit_column(*args, **kwargs) - dot_count = result.count('.') - assert dot_count in (0, 1, 2), "Unexpected visit_column result {}".format(result) - if dot_count == 2: - # we have something of the form schema.table.column - # hive doesn't like the schema in front, so chop it out - result = result[result.index('.') + 1:] - return result - - def visit_char_length_func(self, fn, **kw): - return 'length{}'.format(self.function_argspec(fn, **kw)) - - -class HiveTypeCompiler(compiler.GenericTypeCompiler): - def visit_INTEGER(self, type_): - return 'INT' - - def visit_NUMERIC(self, type_): - return 'DECIMAL' - - def visit_CHAR(self, type_): - return 'STRING' - - def visit_VARCHAR(self, type_): - return 'STRING' - - def visit_NCHAR(self, type_): - return 'STRING' - - def visit_TEXT(self, type_): - return 'STRING' - - def visit_CLOB(self, type_): - return 'STRING' - - def visit_BLOB(self, type_): - return 'BINARY' - - def visit_TIME(self, type_): - return 'TIMESTAMP' - - def visit_DATE(self, type_): - return 'TIMESTAMP' - - def visit_DATETIME(self, type_): - return 'TIMESTAMP' - - -class HiveExecutionContext(default.DefaultExecutionContext): - """This is pretty much the same as SQLiteExecutionContext to work around the same issue. - - http://docs.sqlalchemy.org/en/latest/dialects/sqlite.html#dotted-column-names - - engine = create_engine('hive://...', execution_options={'hive_raw_colnames': True}) - """ - - @util.memoized_property - def _preserve_raw_colnames(self): - # Ideally, this would also gate on hive.resultset.use.unique.column.names - return self.execution_options.get('hive_raw_colnames', False) - - def _translate_colname(self, colname): - # Adjust for dotted column names. - # When hive.resultset.use.unique.column.names is true (the default), Hive returns column - # names as "tablename.colname" in cursor.description. - if not self._preserve_raw_colnames and '.' in colname: - return colname.split('.')[-1], colname - else: - return colname, None - - -class HiveDialect(default.DefaultDialect): - name = 'hive' - driver = 'thrift' - execution_ctx_cls = HiveExecutionContext - preparer = HiveIdentifierPreparer - statement_compiler = HiveCompiler - supports_views = True - supports_alter = True - supports_pk_autoincrement = False - supports_default_values = False - supports_empty_insert = False - supports_native_decimal = True - supports_native_boolean = True - supports_unicode_statements = True - supports_unicode_binds = True - returns_unicode_strings = True - description_encoding = None - supports_multivalues_insert = True - type_compiler = HiveTypeCompiler - supports_sane_rowcount = False - supports_statement_cache = False - - @classmethod - def dbapi(cls): - return hive - - @classmethod - def import_dbapi(cls): - return hive - - def create_connect_args(self, url): - kwargs = { - 'host': url.host, - 'port': url.port or 10000, - 'username': url.username, - 'password': url.password, - 'database': url.database or 'default', - } - kwargs.update(url.query) - return [], kwargs - - def get_schema_names(self, connection, **kw): - # Equivalent to SHOW DATABASES - return [row[0] for row in connection.execute(text('SHOW SCHEMAS'))] - - def get_view_names(self, connection, schema=None, **kw): - # Hive does not provide functionality to query tableType - # This allows reflection to not crash at the cost of being inaccurate - return self.get_table_names(connection, schema, **kw) - - def _get_table_columns(self, connection, table_name, schema): - full_table = table_name - if schema: - full_table = schema + '.' + table_name - # TODO using TGetColumnsReq hangs after sending TFetchResultsReq. - # Using DESCRIBE works but is uglier. - try: - # This needs the table name to be unescaped (no backticks). - rows = connection.execute(text('DESCRIBE {}'.format(full_table))).fetchall() - except exc.OperationalError as e: - # Does the table exist? - regex_fmt = r'TExecuteStatementResp.*SemanticException.*Table not found {}' - regex = regex_fmt.format(re.escape(full_table)) - if re.search(regex, e.args[0]): - raise exc.NoSuchTableError(full_table) - else: - raise - else: - # Hive is stupid: this is what I get from DESCRIBE some_schema.does_not_exist - regex = r'Table .* does not exist' - if len(rows) == 1 and re.match(regex, rows[0].col_name): - raise exc.NoSuchTableError(full_table) - return rows - - def has_table(self, connection, table_name, schema=None, **kw): - try: - self._get_table_columns(connection, table_name, schema) - return True - except exc.NoSuchTableError: - return False - - def get_columns(self, connection, table_name, schema=None, **kw): - rows = self._get_table_columns(connection, table_name, schema) - # Strip whitespace - rows = [[col.strip() if col else None for col in row] for row in rows] - # Filter out empty rows and comment - rows = [row for row in rows if row[0] and row[0] != '# col_name'] - result = [] - for (col_name, col_type, _comment) in rows: - if col_name == '# Partition Information': - break - # Take out the more detailed type information - # e.g. 'map' -> 'map' - # 'decimal(10,1)' -> decimal - col_type = re.search(r'^\w+', col_type).group(0) - try: - coltype = _type_map[col_type] - except KeyError: - util.warn("Did not recognize type '%s' of column '%s'" % (col_type, col_name)) - coltype = types.NullType - - result.append({ - 'name': col_name, - 'type': coltype, - 'nullable': True, - 'default': None, - }) - return result - - def get_foreign_keys(self, connection, table_name, schema=None, **kw): - # Hive has no support for foreign keys. - return [] - - def get_pk_constraint(self, connection, table_name, schema=None, **kw): - # Hive has no support for primary keys. - return [] - - def get_indexes(self, connection, table_name, schema=None, **kw): - rows = self._get_table_columns(connection, table_name, schema) - # Strip whitespace - rows = [[col.strip() if col else None for col in row] for row in rows] - # Filter out empty rows and comment - rows = [row for row in rows if row[0] and row[0] != '# col_name'] - for i, (col_name, _col_type, _comment) in enumerate(rows): - if col_name == '# Partition Information': - break - # Handle partition columns - col_names = [] - for col_name, _col_type, _comment in rows[i + 1:]: - col_names.append(col_name) - if col_names: - return [{'name': 'partition', 'column_names': col_names, 'unique': False}] - else: - return [] - - def get_table_names(self, connection, schema=None, **kw): - query = 'SHOW TABLES' - if schema: - query += ' IN ' + self.identifier_preparer.quote_identifier(schema) - - table_names = [] - - for row in connection.execute(text(query)): - # Hive returns 1 columns - if len(row) == 1: - table_names.append(row[0]) - # Spark SQL returns 3 columns - elif len(row) == 3: - table_names.append(row[1]) - else: - _logger.warning("Unexpected number of columns in SHOW TABLES result: {}".format(len(row))) - table_names.append('UNKNOWN') - - return table_names - - def do_rollback(self, dbapi_connection): - # No transactions for Hive - pass - - def _check_unicode_returns(self, connection, additional_tests=None): - # We decode everything as UTF-8 - return True - - def _check_unicode_description(self, connection): - # We decode everything as UTF-8 - return True - - -class HiveHTTPDialect(HiveDialect): - - name = "hive" - scheme = "http" - driver = "rest" - - def create_connect_args(self, url): - kwargs = { - "host": url.host, - "port": url.port or 10000, - "scheme": self.scheme, - "username": url.username or None, - "password": url.password or None, - "database": url.database or "default", - } - if url.query: - kwargs.update(url.query) - return [], kwargs - return ([], kwargs) - - -class HiveHTTPSDialect(HiveHTTPDialect): - - name = "hive" - scheme = "https" diff --git a/src/chronify/_vendor/kyuubi/pyhive/sqlalchemy_presto.py b/src/chronify/_vendor/kyuubi/pyhive/sqlalchemy_presto.py deleted file mode 100644 index f5a256f..0000000 --- a/src/chronify/_vendor/kyuubi/pyhive/sqlalchemy_presto.py +++ /dev/null @@ -1,256 +0,0 @@ -"""Integration between SQLAlchemy and Presto. - -Some code based on -https://github.com/zzzeek/sqlalchemy/blob/rel_0_5/lib/sqlalchemy/databases/sqlite.py -which is released under the MIT license. -""" - -from __future__ import absolute_import -from __future__ import unicode_literals - -import re -import sqlalchemy -from sqlalchemy import exc -from sqlalchemy import types -from sqlalchemy import util -# TODO shouldn't use mysql type -from sqlalchemy.sql import text -try: - from sqlalchemy.databases import mysql - mysql_tinyinteger = mysql.MSTinyInteger -except ImportError: - # Required for SQLAlchemy>=2.0 - from sqlalchemy.dialects import mysql - mysql_tinyinteger = mysql.base.MSTinyInteger -from sqlalchemy.engine import default -from sqlalchemy.sql import compiler, bindparam -from sqlalchemy.sql.compiler import SQLCompiler - -from pyhive import presto -from pyhive.common import UniversalSet - -sqlalchemy_version = float(re.search(r"^([\d]+\.[\d]+)\..+", sqlalchemy.__version__).group(1)) - -class PrestoIdentifierPreparer(compiler.IdentifierPreparer): - # Just quote everything to make things simpler / easier to upgrade - reserved_words = UniversalSet() - - -_type_map = { - 'boolean': types.Boolean, - 'tinyint': mysql_tinyinteger, - 'smallint': types.SmallInteger, - 'integer': types.Integer, - 'bigint': types.BigInteger, - 'real': types.Float, - 'double': types.Float, - 'varchar': types.String, - 'timestamp': types.TIMESTAMP, - 'date': types.DATE, - 'varbinary': types.VARBINARY, -} - - -class PrestoCompiler(SQLCompiler): - def visit_char_length_func(self, fn, **kw): - return 'length{}'.format(self.function_argspec(fn, **kw)) - - -class PrestoTypeCompiler(compiler.GenericTypeCompiler): - def visit_CLOB(self, type_, **kw): - raise ValueError("Presto does not support the CLOB column type.") - - def visit_NCLOB(self, type_, **kw): - raise ValueError("Presto does not support the NCLOB column type.") - - def visit_DATETIME(self, type_, **kw): - raise ValueError("Presto does not support the DATETIME column type.") - - def visit_FLOAT(self, type_, **kw): - return 'DOUBLE' - - def visit_TEXT(self, type_, **kw): - if type_.length: - return 'VARCHAR({:d})'.format(type_.length) - else: - return 'VARCHAR' - - -class PrestoDialect(default.DefaultDialect): - name = 'presto' - driver = 'rest' - paramstyle = 'pyformat' - preparer = PrestoIdentifierPreparer - statement_compiler = PrestoCompiler - supports_alter = False - supports_pk_autoincrement = False - supports_default_values = False - supports_empty_insert = False - supports_multivalues_insert = True - supports_unicode_statements = True - supports_unicode_binds = True - supports_statement_cache = False - returns_unicode_strings = True - description_encoding = None - supports_native_boolean = True - type_compiler = PrestoTypeCompiler - - @classmethod - def dbapi(cls): - return presto - - @classmethod - def import_dbapi(cls): - return presto - - def create_connect_args(self, url): - db_parts = (url.database or 'hive').split('/') - kwargs = { - 'host': url.host, - 'port': url.port or 8080, - 'username': url.username, - 'password': url.password - } - kwargs.update(url.query) - if len(db_parts) == 1: - kwargs['catalog'] = db_parts[0] - elif len(db_parts) == 2: - kwargs['catalog'] = db_parts[0] - kwargs['schema'] = db_parts[1] - else: - raise ValueError("Unexpected database format {}".format(url.database)) - return [], kwargs - - def get_schema_names(self, connection, **kw): - return [row.Schema for row in connection.execute(text('SHOW SCHEMAS'))] - - def _get_table_columns(self, connection, table_name, schema): - full_table = self.identifier_preparer.quote_identifier(table_name) - if schema: - full_table = self.identifier_preparer.quote_identifier(schema) + '.' + full_table - try: - return connection.execute(text('SHOW COLUMNS FROM {}'.format(full_table))) - except (presto.DatabaseError, exc.DatabaseError) as e: - # Normally SQLAlchemy should wrap this exception in sqlalchemy.exc.DatabaseError, which - # it successfully does in the Hive version. The difference with Presto is that this - # error is raised when fetching the cursor's description rather than the initial execute - # call. SQLAlchemy doesn't handle this. Thus, we catch the unwrapped - # presto.DatabaseError here. - # Does the table exist? - msg = ( - e.args[0].get('message') if e.args and isinstance(e.args[0], dict) - else e.args[0] if e.args and isinstance(e.args[0], str) - else None - ) - regex = r"Table\ \'.*{}\'\ does\ not\ exist".format(re.escape(table_name)) - if msg and re.search(regex, msg, re.IGNORECASE): - raise exc.NoSuchTableError(table_name) - else: - raise - - def has_table(self, connection, table_name, schema=None, **kw): - try: - self._get_table_columns(connection, table_name, schema) - return True - except exc.NoSuchTableError: - return False - - def get_columns(self, connection, table_name, schema=None, **kw): - rows = self._get_table_columns(connection, table_name, schema) - result = [] - for row in rows: - try: - coltype = _type_map[row.Type] - except KeyError: - util.warn("Did not recognize type '%s' of column '%s'" % (row.Type, row.Column)) - coltype = types.NullType - result.append({ - 'name': row.Column, - 'type': coltype, - # newer Presto no longer includes this column - 'nullable': getattr(row, 'Null', True), - 'default': None, - }) - return result - - def get_foreign_keys(self, connection, table_name, schema=None, **kw): - # Hive has no support for foreign keys. - return [] - - def get_pk_constraint(self, connection, table_name, schema=None, **kw): - # Hive has no support for primary keys. - return [] - - def get_indexes(self, connection, table_name, schema=None, **kw): - rows = self._get_table_columns(connection, table_name, schema) - col_names = [] - for row in rows: - part_key = 'Partition Key' - # Presto puts this information in one of 3 places depending on version - # - a boolean column named "Partition Key" - # - a string in the "Comment" column - # - a string in the "Extra" column - if sqlalchemy_version >= 1.4: - row = row._mapping - is_partition_key = ( - (part_key in row and row[part_key]) - or row['Comment'].startswith(part_key) - or ('Extra' in row and 'partition key' in row['Extra']) - ) - if is_partition_key: - col_names.append(row['Column']) - if col_names: - return [{'name': 'partition', 'column_names': col_names, 'unique': False}] - else: - return [] - - def _get_default_schema_name(self, connection): - #'SELECT CURRENT_SCHEMA()' - return super()._get_default_schema_name(connection) - - def get_table_names(self, connection, schema=None, **kw): - query = 'SHOW TABLES' - # N.B. This is incorrect, if no schema is provided, the current/default schema should be used - # with a call to an overridden self._get_default_schema_name(connection), but I could not - # see how to implement that as there is no CURRENT_SCHEMA function - # default_schema = self._get_default_schema_name(connection) - - if schema: - query += ' FROM ' + self.identifier_preparer.quote_identifier(schema) - return [row.Table for row in connection.execute(text(query))] - - def get_view_names(self, connection, schema=None, **kw): - if schema: - view_name_query = """ - SELECT table_name - FROM information_schema.views - WHERE table_schema = :schema - """ - query = text(view_name_query).bindparams( - bindparam("schema", type_=types.Unicode) - ) - else: - # N.B. This is incorrect, if no schema is provided, the current/default schema should - # be used with a call to self._get_default_schema_name(connection), but I could not - # see how to implement that - # default_schema = self._get_default_schema_name(connection) - view_name_query = """ - SELECT table_name - FROM information_schema.views - """ - query = text(view_name_query) - - result = connection.execute(query, dict(schema=schema)) - return [row[0] for row in result] - - def do_rollback(self, dbapi_connection): - # No transactions for Presto - pass - - def _check_unicode_returns(self, connection, additional_tests=None): - # requests gives back Unicode strings - return True - - def _check_unicode_description(self, connection): - # requests gives back Unicode strings - return True diff --git a/src/chronify/_vendor/kyuubi/pyhive/sqlalchemy_trino.py b/src/chronify/_vendor/kyuubi/pyhive/sqlalchemy_trino.py deleted file mode 100644 index 11be2a6..0000000 --- a/src/chronify/_vendor/kyuubi/pyhive/sqlalchemy_trino.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Integration between SQLAlchemy and Trino. - -Some code based on -https://github.com/zzzeek/sqlalchemy/blob/rel_0_5/lib/sqlalchemy/databases/sqlite.py -which is released under the MIT license. -""" - -from __future__ import absolute_import -from __future__ import unicode_literals - -import re -from sqlalchemy import exc -from sqlalchemy import types -from sqlalchemy import util -# TODO shouldn't use mysql type -try: - from sqlalchemy.databases import mysql - mysql_tinyinteger = mysql.MSTinyInteger -except ImportError: - # Required for SQLAlchemy>=2.0 - from sqlalchemy.dialects import mysql - mysql_tinyinteger = mysql.base.MSTinyInteger -from sqlalchemy.engine import default -from sqlalchemy.sql import compiler -from sqlalchemy.sql.compiler import SQLCompiler - -from pyhive import trino -from pyhive.common import UniversalSet -from pyhive.sqlalchemy_presto import PrestoDialect, PrestoCompiler, PrestoIdentifierPreparer - -class TrinoIdentifierPreparer(PrestoIdentifierPreparer): - pass - - -_type_map = { - 'boolean': types.Boolean, - 'tinyint': mysql_tinyinteger, - 'smallint': types.SmallInteger, - 'integer': types.Integer, - 'bigint': types.BigInteger, - 'real': types.Float, - 'double': types.Float, - 'varchar': types.String, - 'timestamp': types.TIMESTAMP, - 'date': types.DATE, - 'varbinary': types.VARBINARY, -} - - -class TrinoCompiler(PrestoCompiler): - pass - - -class TrinoTypeCompiler(PrestoCompiler): - def visit_CLOB(self, type_, **kw): - raise ValueError("Trino does not support the CLOB column type.") - - def visit_NCLOB(self, type_, **kw): - raise ValueError("Trino does not support the NCLOB column type.") - - def visit_DATETIME(self, type_, **kw): - raise ValueError("Trino does not support the DATETIME column type.") - - def visit_FLOAT(self, type_, **kw): - return 'DOUBLE' - - def visit_TEXT(self, type_, **kw): - if type_.length: - return 'VARCHAR({:d})'.format(type_.length) - else: - return 'VARCHAR' - - -class TrinoDialect(PrestoDialect): - name = 'trino' - supports_statement_cache = False - - @classmethod - def dbapi(cls): - return trino - - @classmethod - def import_dbapi(cls): - return trino diff --git a/src/chronify/_vendor/kyuubi/pyhive/trino.py b/src/chronify/_vendor/kyuubi/pyhive/trino.py deleted file mode 100644 index 658457a..0000000 --- a/src/chronify/_vendor/kyuubi/pyhive/trino.py +++ /dev/null @@ -1,144 +0,0 @@ -"""DB-API implementation backed by Trino - -See http://www.python.org/dev/peps/pep-0249/ - -Many docstrings in this file are based on the PEP, which is in the public domain. -""" - -from __future__ import absolute_import -from __future__ import unicode_literals - -import logging - -import requests - -# Make all exceptions visible in this module per DB-API -from pyhive.common import DBAPITypeObject -from pyhive.exc import * # noqa -from pyhive.presto import Connection as PrestoConnection, Cursor as PrestoCursor, PrestoParamEscaper - -try: # Python 3 - import urllib.parse as urlparse -except ImportError: # Python 2 - import urlparse - -# PEP 249 module globals -apilevel = '2.0' -threadsafety = 2 # Threads may share the module and connections. -paramstyle = 'pyformat' # Python extended format codes, e.g. ...WHERE name=%(name)s - -_logger = logging.getLogger(__name__) - - -class TrinoParamEscaper(PrestoParamEscaper): - pass - - -_escaper = TrinoParamEscaper() - - -def connect(*args, **kwargs): - """Constructor for creating a connection to the database. See class :py:class:`Connection` for - arguments. - - :returns: a :py:class:`Connection` object. - """ - return Connection(*args, **kwargs) - - -class Connection(PrestoConnection): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def cursor(self): - """Return a new :py:class:`Cursor` object using the connection.""" - return Cursor(*self._args, **self._kwargs) - - -class Cursor(PrestoCursor): - """These objects represent a database cursor, which is used to manage the context of a fetch - operation. - - Cursors are not isolated, i.e., any changes done to the database by a cursor are immediately - visible by other cursors or connections. - """ - - def execute(self, operation, parameters=None): - """Prepare and execute a database operation (query or command). - - Return values are not defined. - """ - headers = { - 'X-Trino-Catalog': self._catalog, - 'X-Trino-Schema': self._schema, - 'X-Trino-Source': self._source, - 'X-Trino-User': self._username, - } - - if self._session_props: - headers['X-Trino-Session'] = ','.join( - '{}={}'.format(propname, propval) - for propname, propval in self._session_props.items() - ) - - # Prepare statement - if parameters is None: - sql = operation - else: - sql = operation % _escaper.escape_args(parameters) - - self._reset_state() - - self._state = self._STATE_RUNNING - url = urlparse.urlunparse(( - self._protocol, - '{}:{}'.format(self._host, self._port), '/v1/statement', None, None, None)) - _logger.info('%s', sql) - _logger.debug("Headers: %s", headers) - response = self._requests_session.post( - url, data=sql.encode('utf-8'), headers=headers, **self._requests_kwargs) - self._process_response(response) - - def _process_response(self, response): - """Given the JSON response from Trino's REST API, update the internal state with the next - URI and any data from the response - """ - # TODO handle HTTP 503 - if response.status_code != requests.codes.ok: - fmt = "Unexpected status code {}\n{}" - raise OperationalError(fmt.format(response.status_code, response.content)) - - response_json = response.json() - _logger.debug("Got response %s", response_json) - assert self._state == self._STATE_RUNNING, "Should be running if processing response" - self._nextUri = response_json.get('nextUri') - self._columns = response_json.get('columns') - if 'id' in response_json: - self.last_query_id = response_json['id'] - if 'X-Trino-Clear-Session' in response.headers: - propname = response.headers['X-Trino-Clear-Session'] - self._session_props.pop(propname, None) - if 'X-Trino-Set-Session' in response.headers: - propname, propval = response.headers['X-Trino-Set-Session'].split('=', 1) - self._session_props[propname] = propval - if 'data' in response_json: - assert self._columns - new_data = response_json['data'] - self._process_data(new_data) - self._data += map(tuple, new_data) - if 'nextUri' not in response_json: - self._state = self._STATE_FINISHED - if 'error' in response_json: - raise DatabaseError(response_json['error']) - - -# -# Type Objects and Constructors -# - - -# See types in trino-main/src/main/java/com/facebook/trino/tuple/TupleInfo.java -FIXED_INT_64 = DBAPITypeObject(['bigint']) -VARIABLE_BINARY = DBAPITypeObject(['varchar']) -DOUBLE = DBAPITypeObject(['double']) -BOOLEAN = DBAPITypeObject(['boolean']) diff --git a/src/chronify/csv_io.py b/src/chronify/csv_io.py index 5b99b07..38a7159 100644 --- a/src/chronify/csv_io.py +++ b/src/chronify/csv_io.py @@ -4,16 +4,15 @@ import duckdb from duckdb import DuckDBPyRelation -from chronify.models import CsvTableSchema, get_duckdb_type_from_sqlalchemy +from chronify.models import CsvTableSchema +from chronify.ibis.types import get_duckdb_type_from_ibis from chronify.time_configs import DatetimeRange def read_csv(path: Path | str, schema: CsvTableSchema, **kwargs: Any) -> DuckDBPyRelation: """Read a CSV file into a DuckDB relation.""" if schema.column_dtypes: - dtypes = { - x.name: get_duckdb_type_from_sqlalchemy(x.dtype).id for x in schema.column_dtypes - } + dtypes = {x.name: get_duckdb_type_from_ibis(x.dtype).id for x in schema.column_dtypes} rel = duckdb.read_csv(str(path), dtype=dtypes, **kwargs) else: rel = duckdb.read_csv(str(path), **kwargs) diff --git a/src/chronify/ibis/__init__.py b/src/chronify/ibis/__init__.py new file mode 100644 index 0000000..5f4d7e0 --- /dev/null +++ b/src/chronify/ibis/__init__.py @@ -0,0 +1,19 @@ +"""Ibis-based database backend abstraction layer.""" + +from chronify.ibis.backend import ( + DuckDBBackend, + IbisBackend, + SparkBackend, + SQLiteBackend, + get_backend, + make_backend, +) + +__all__ = [ + "DuckDBBackend", + "IbisBackend", + "SparkBackend", + "SQLiteBackend", + "get_backend", + "make_backend", +] diff --git a/src/chronify/ibis/backend.py b/src/chronify/ibis/backend.py new file mode 100644 index 0000000..9bab6ac --- /dev/null +++ b/src/chronify/ibis/backend.py @@ -0,0 +1,73 @@ +"""Ibis backend abstraction layer for database operations. + +This module provides a unified interface for different database backends +(DuckDB, SQLite, Spark) using the Ibis library. + +The module re-exports all backend classes and factory functions for convenience. +Individual backend implementations are in separate modules: +- base.py: IbisBackend abstract base class +- duckdb_backend.py: DuckDBBackend +- sqlite_backend.py: SQLiteBackend +- spark_backend.py: SparkBackend +""" + +from pathlib import Path +from typing import Any + +from chronify.ibis.base import IbisBackend +from chronify.ibis.duckdb_backend import DuckDBBackend +from chronify.ibis.sqlite_backend import SQLiteBackend +from chronify.ibis.spark_backend import SparkBackend + + +def make_backend( + backend_name: str = "duckdb", + file_path: Path | str | None = None, + **kwargs: Any, +) -> IbisBackend: + """Factory function to create an IbisBackend instance. + + Parameters + ---------- + backend_name + Name of the backend: "duckdb", "sqlite", or "spark" + file_path + Path to database file (for duckdb/sqlite), or None for in-memory + **kwargs + Additional arguments passed to the backend constructor + + Returns + ------- + IbisBackend + Backend instance + + Examples + -------- + >>> backend = make_backend("duckdb") # In-memory DuckDB + >>> backend = make_backend("duckdb", file_path="data.db") # File-based + >>> backend = make_backend("sqlite") # In-memory SQLite + >>> backend = make_backend("spark") # Spark (requires spark extra) + """ + match backend_name: + case "duckdb": + return DuckDBBackend(file_path=file_path, **kwargs) + case "sqlite": + return SQLiteBackend(file_path=file_path, **kwargs) + case "spark": + return SparkBackend(**kwargs) + case _: + msg = f"Unknown backend: {backend_name}. Supported: duckdb, sqlite, spark" + raise NotImplementedError(msg) + + +# Backwards compatibility alias +get_backend = make_backend + +__all__ = [ + "IbisBackend", + "DuckDBBackend", + "SQLiteBackend", + "SparkBackend", + "make_backend", + "get_backend", +] diff --git a/src/chronify/ibis/base.py b/src/chronify/ibis/base.py new file mode 100644 index 0000000..740cdb2 --- /dev/null +++ b/src/chronify/ibis/base.py @@ -0,0 +1,280 @@ +"""Abstract base class for Ibis-based database backends.""" + +from abc import ABC, abstractmethod +from contextlib import contextmanager +from pathlib import Path +from typing import Iterator, Optional, cast + +import ibis +import ibis.expr.types as ir +import pandas as pd +import pyarrow as pa +from loguru import logger + + +class IbisBackend(ABC): + """Abstract base class for Ibis-based database backends.""" + + def __init__(self, connection: ibis.BaseBackend) -> None: + self._connection = connection + + @property + def connection(self) -> ibis.BaseBackend: + """Return the underlying Ibis connection.""" + return self._connection + + @property + @abstractmethod + def name(self) -> str: + """Return the name of the backend (e.g., 'duckdb', 'sqlite', 'spark').""" + + @property + @abstractmethod + def database(self) -> str | None: + """Return the database path/name, or None for in-memory.""" + + @abstractmethod + def create_table( + self, + name: str, + data: pd.DataFrame | pa.Table | ir.Table, + overwrite: bool = False, + ) -> ir.Table: + """Create a table from data. + + Parameters + ---------- + name + Name of the table to create + data + Data to populate the table + overwrite + If True, drop existing table first + + Returns + ------- + ir.Table + Ibis table reference + """ + + @abstractmethod + def create_view(self, name: str, expr: ir.Table) -> None: + """Create a view from an Ibis expression. + + Parameters + ---------- + name + Name of the view + expr + Ibis table expression defining the view + """ + + @abstractmethod + def create_temp_view(self, name: str, data: ir.Table | pd.DataFrame | pa.Table) -> None: + """Create a temporary view from user-provided data. + + Parameters + ---------- + name + Name of the temporary view + data + Data for the view + """ + + @abstractmethod + def drop_table(self, name: str, if_exists: bool = False) -> None: + """Drop a table. + + Parameters + ---------- + name + Name of the table to drop + if_exists + If True, don't raise error if table doesn't exist + """ + + @abstractmethod + def drop_view(self, name: str, if_exists: bool = False) -> None: + """Drop a view. + + Parameters + ---------- + name + Name of the view to drop + if_exists + If True, don't raise error if view doesn't exist + """ + + @abstractmethod + def list_tables(self) -> list[str]: + """Return a list of all tables and views in the database.""" + + @abstractmethod + def table(self, name: str) -> ir.Table: + """Get a reference to an existing table. + + Parameters + ---------- + name + Name of the table + + Returns + ------- + ir.Table + Ibis table reference + """ + + @abstractmethod + def execute(self, expr: ir.Expr) -> pd.DataFrame: + """Execute an Ibis expression and return results as DataFrame. + + Parameters + ---------- + expr + Ibis expression to execute + + Returns + ------- + pd.DataFrame + Query results + """ + + @abstractmethod + def insert(self, table_name: str, data: pd.DataFrame | pa.Table | ir.Table) -> None: + """Insert data into an existing table. + + Parameters + ---------- + table_name + Name of the target table + data + Data to insert + """ + + @abstractmethod + def sql(self, query: str) -> ir.Table: + """Execute raw SQL and return an Ibis table expression. + + Parameters + ---------- + query + SQL query string + + Returns + ------- + ir.Table + Query results as Ibis table + """ + + def execute_sql(self, query: str) -> None: + """Execute raw SQL without returning results. + + Parameters + ---------- + query + SQL statement to execute + """ + self._connection.raw_sql(query) + + def execute_sql_to_df(self, query: str) -> pd.DataFrame: + """Execute raw SQL and return results as a DataFrame. + + This method is useful for complex queries that may conflict with Ibis's + SQL parsing, such as queries with CTEs. + + Parameters + ---------- + query + SQL query string + + Returns + ------- + pd.DataFrame + Query results + """ + return cast(pd.DataFrame, self._connection.raw_sql(query).fetchdf()) + + @abstractmethod + def write_parquet( + self, + expr: ir.Table, + path: Path, + partition_by: Optional[list[str]] = None, + overwrite: bool = False, + ) -> None: + """Write query results to a Parquet file. + + Parameters + ---------- + expr + Ibis table expression to write + path + Output path + partition_by + Optional columns to partition by + overwrite + If True, overwrite existing file/directory + """ + + @abstractmethod + def create_view_from_parquet(self, name: str, path: Path) -> None: + """Create a view that reads from a Parquet file. + + Parameters + ---------- + name + Name of the view to create + path + Path to Parquet file or directory + """ + + def has_table(self, name: str) -> bool: + """Check if a table or view exists. + + Parameters + ---------- + name + Name of the table/view + + Returns + ------- + bool + True if exists + """ + return name in self.list_tables() + + def dispose(self) -> None: + """Clean up the backend connection.""" + pass + + def reconnect(self) -> None: + """Reconnect to the database after dispose.""" + pass + + @contextmanager + def transaction(self) -> Iterator[list[tuple[str, str]]]: + """Context manager for pseudo-transaction support. + + Ibis lacks built-in transaction support, so we track created objects + and clean them up on failure. + + Yields + ------ + list[tuple[str, str]] + List to track created objects as (type, name) tuples. + Callers should append to this list when creating tables/views. + """ + created_objects: list[tuple[str, str]] = [] + try: + yield created_objects + except Exception: + # Rollback by dropping created objects in reverse order + for obj_type, name in reversed(created_objects): + try: + if obj_type == "table": + self.drop_table(name, if_exists=True) + elif obj_type == "view": + self.drop_view(name, if_exists=True) + except Exception as e: + logger.warning("Failed to rollback {} {}: {}", obj_type, name, e) + raise diff --git a/src/chronify/ibis/duckdb_backend.py b/src/chronify/ibis/duckdb_backend.py new file mode 100644 index 0000000..949722b --- /dev/null +++ b/src/chronify/ibis/duckdb_backend.py @@ -0,0 +1,157 @@ +"""DuckDB backend implementation.""" + +from pathlib import Path +from typing import Any, Optional, cast + +import ibis +import ibis.expr.types as ir +import pandas as pd +import pyarrow as pa + +from chronify.exceptions import InvalidOperation +from chronify.ibis.base import IbisBackend + + +class DuckDBBackend(IbisBackend): + """DuckDB backend implementation.""" + + def __init__( + self, + file_path: Optional[Path | str] = None, + **kwargs: Any, + ) -> None: + """Create a DuckDB backend. + + Parameters + ---------- + file_path + Path to database file, or None for in-memory + **kwargs + Additional arguments passed to ibis.duckdb.connect + """ + self._file_path = str(file_path) if file_path else ":memory:" + connection = ibis.duckdb.connect( + self._file_path, + **kwargs, + ) + super().__init__(connection) + + @property + def name(self) -> str: + return "duckdb" + + @property + def database(self) -> str | None: + return None if self._file_path == ":memory:" else self._file_path + + def create_table( + self, + name: str, + data: pd.DataFrame | pa.Table | ir.Table, + overwrite: bool = False, + ) -> ir.Table: + if overwrite and self.has_table(name): + self.drop_table(name, if_exists=True) + + if isinstance(data, ir.Table): + # Execute the expression and create table from result + self._connection.raw_sql(f"CREATE TABLE {name} AS {data.compile()}") + else: + self._connection.create_table(name, data) + + return self.table(name) + + def create_view(self, name: str, expr: ir.Table) -> None: + sql = expr.compile() + self._connection.raw_sql(f"CREATE VIEW {name} AS {sql}") + + def create_temp_view(self, name: str, data: ir.Table | pd.DataFrame | pa.Table) -> None: + if isinstance(data, ir.Table): + sql = data.compile() + self._connection.raw_sql(f"CREATE OR REPLACE TEMP VIEW {name} AS {sql}") + else: + # Register as a table then create view + temp_name = f"_temp_{name}" + self._connection.create_table(temp_name, data, temp=True) + self._connection.raw_sql( + f"CREATE OR REPLACE TEMP VIEW {name} AS SELECT * FROM {temp_name}" + ) + + def drop_table(self, name: str, if_exists: bool = False) -> None: + if_exists_sql = "IF EXISTS " if if_exists else "" + self._connection.raw_sql(f"DROP TABLE {if_exists_sql}{name}") + + def drop_view(self, name: str, if_exists: bool = False) -> None: + if_exists_sql = "IF EXISTS " if if_exists else "" + self._connection.raw_sql(f"DROP VIEW {if_exists_sql}{name}") + + def list_tables(self) -> list[str]: + # Filter out internal Ibis memtables + return [ + t + for t in self._connection.list_tables() + if not t.startswith("ibis_pandas_memtable_") + and not t.startswith("ibis_pyarrow_memtable_") + and not t.startswith("_temp_") + and not t.startswith("_backing_") + ] + + def table(self, name: str) -> ir.Table: + return self._connection.table(name) + + def execute(self, expr: ir.Expr) -> pd.DataFrame: + return cast(pd.DataFrame, expr.execute()) + + def insert(self, table_name: str, data: pd.DataFrame | pa.Table | ir.Table) -> None: + if isinstance(data, ir.Table): + sql = data.compile() + self._connection.raw_sql(f"INSERT INTO {table_name} {sql}") + elif isinstance(data, pa.Table): + # Register PyArrow table as a virtual table, insert, then unregister + self._connection.con.register("__temp_insert_df", data) + self._connection.con.execute( + f"INSERT INTO {table_name} SELECT * FROM __temp_insert_df" + ) + self._connection.con.unregister("__temp_insert_df") + else: + # pandas DataFrame - register as a virtual table, insert, then unregister + self._connection.con.register("__temp_insert_df", data) + self._connection.con.execute( + f"INSERT INTO {table_name} SELECT * FROM __temp_insert_df" + ) + self._connection.con.unregister("__temp_insert_df") + + def sql(self, query: str) -> ir.Table: + return self._connection.sql(query) + + def write_parquet( + self, + expr: ir.Table, + path: Path, + partition_by: Optional[list[str]] = None, + overwrite: bool = False, + ) -> None: + if not overwrite and path.exists(): + msg = f"File already exists: {path}" + raise InvalidOperation(msg) + + sql = expr.compile() + if partition_by: + cols = ",".join(partition_by) + query = f"COPY ({sql}) TO '{path}' (FORMAT PARQUET, PARTITION_BY ({cols}))" + else: + query = f"COPY ({sql}) TO '{path}' (FORMAT PARQUET)" + self._connection.raw_sql(query) + + def create_view_from_parquet(self, name: str, path: Path) -> None: + str_path = f"{path}/**/*.parquet" if path.is_dir() else str(path) + query = f"CREATE VIEW {name} AS SELECT * FROM read_parquet('{str_path}')" + self._connection.raw_sql(query) + + def dispose(self) -> None: + """Close the DuckDB connection to flush data to disk.""" + self._connection.con.close() + + def reconnect(self) -> None: + """Reconnect to the database after dispose.""" + self._connection = ibis.duckdb.connect(self._file_path) diff --git a/src/chronify/ibis/functions.py b/src/chronify/ibis/functions.py new file mode 100644 index 0000000..11fd0a3 --- /dev/null +++ b/src/chronify/ibis/functions.py @@ -0,0 +1,352 @@ +"""Database I/O functions using Ibis backends. + +This module provides functions to read and write data efficiently using Ibis. +""" + +from collections import Counter +from pathlib import Path +from typing import Sequence + +import ibis.expr.types as ir +import pandas as pd +import pyarrow as pa +from numpy.dtypes import ObjectDType +from pandas import DatetimeTZDtype + +from chronify.exceptions import InvalidOperation, InvalidParameter +from chronify.ibis.backend import IbisBackend +from chronify.time import TimeDataType +from chronify.time_configs import ( + DatetimeRange, + DatetimeRangeBase, + DatetimeRanges, + DatetimeRangeWithTZColumn, + TimeBaseModel, +) +from chronify.utils.path_utils import check_overwrite + +# Tuple for isinstance checks to enable proper type narrowing +_DATETIME_RANGES = (DatetimeRange, DatetimeRangeWithTZColumn) + + +def read_table( + backend: IbisBackend, + table_name: str, + config: TimeBaseModel, +) -> pd.DataFrame: + """Read a table from the database. + + Parameters + ---------- + backend + Ibis backend + table_name + Name of table to read + config + Time configuration for datetime conversion + + Returns + ------- + pd.DataFrame + Table data as DataFrame + """ + table = backend.table(table_name) + df = backend.execute(table) + + if backend.name == "sqlite" and isinstance(config, _DATETIME_RANGES): + _convert_database_output_for_datetime(df, config) + elif backend.name == "spark" and isinstance(config, _DATETIME_RANGES): + _convert_spark_output_for_datetime(df, config) + + return df + + +def read_query( + backend: IbisBackend, + expr: ir.Table, + config: TimeBaseModel, +) -> pd.DataFrame: + """Execute an Ibis expression and return results. + + Parameters + ---------- + backend + Ibis backend + expr + Ibis table expression to execute + config + Time configuration for datetime conversion + + Returns + ------- + pd.DataFrame + Query results as DataFrame + """ + df = backend.execute(expr) + + if backend.name == "sqlite" and isinstance(config, _DATETIME_RANGES): + _convert_database_output_for_datetime(df, config) + elif backend.name == "spark" and isinstance(config, _DATETIME_RANGES): + _convert_spark_output_for_datetime(df, config) + + return df + + +def write_table( + backend: IbisBackend, + df: pd.DataFrame | pa.Table, + table_name: str, + configs: Sequence[TimeBaseModel], + if_exists: str = "append", + scratch_dir: Path | None = None, +) -> None: + """Write a DataFrame to the database. + + Parameters + ---------- + backend + Ibis backend + df + DataFrame or PyArrow Table to write + table_name + Name of target table + configs + Time configurations for datetime handling + if_exists + What to do if table exists: "append", "replace", or "fail" + scratch_dir + Directory for temporary files (used by some backends) + """ + match backend.name: + case "duckdb": + _write_to_duckdb(backend, df, table_name, if_exists) + case "sqlite": + _write_to_sqlite(backend, df, table_name, configs, if_exists) + case "spark": + _write_to_pyspark(backend, df, table_name, configs, if_exists, scratch_dir) + case _: + msg = f"Unsupported backend: {backend.name}" + raise NotImplementedError(msg) + + +def write_parquet( + backend: IbisBackend, + query: str | ir.Table, + output_file: Path, + overwrite: bool = False, + partition_columns: list[str] | None = None, +) -> None: + """Write query results to a Parquet file. + + Parameters + ---------- + backend + Ibis backend + query + SQL query string or Ibis table expression + output_file + Output file path + overwrite + If True, overwrite existing file + partition_columns + Optional columns to partition by + """ + check_overwrite(output_file, overwrite) + + if isinstance(query, str): + expr = backend.sql(query) + else: + expr = query + + backend.write_parquet( + expr, + output_file, + partition_by=partition_columns, + overwrite=overwrite, + ) + + +def create_view_from_parquet( + backend: IbisBackend, + view_name: str, + filename: Path, +) -> None: + """Create a view from a Parquet file. + + Parameters + ---------- + backend + Ibis backend + view_name + Name of view to create + filename + Path to Parquet file or directory + """ + backend.create_view_from_parquet(view_name, filename) + + +def _check_one_config_per_datetime_column(configs: Sequence[TimeBaseModel]) -> None: + """Ensure each datetime column has at most one config.""" + time_col_count = Counter( + config.time_column for config in configs if isinstance(config, DatetimeRangeBase) + ) + time_col_dup = {k: v for k, v in time_col_count.items() if v > 1} + if time_col_dup: + msg = f"More than one datetime config found for: {time_col_dup}" + raise InvalidParameter(msg) + + +def _convert_database_input_for_datetime( + df: pd.DataFrame, config: DatetimeRanges, copied: bool +) -> tuple[pd.DataFrame, bool]: + """Convert DataFrame datetime columns for SQLite input.""" + if config.dtype == TimeDataType.TIMESTAMP_NTZ: + return df, copied + + if copied: + df2 = df + else: + df2 = df.copy() + copied = True + + if isinstance(df2[config.time_column].dtype, DatetimeTZDtype): + df2[config.time_column] = df2[config.time_column].dt.tz_convert("UTC") + else: + df2[config.time_column] = df2[config.time_column].dt.tz_localize("UTC") + + return df2, copied + + +def _convert_database_output_for_datetime(df: pd.DataFrame, config: DatetimeRanges) -> None: + """Convert DataFrame datetime columns after SQLite output.""" + if config.time_column in df.columns: + if config.dtype == TimeDataType.TIMESTAMP_TZ: + if isinstance(df[config.time_column].dtype, ObjectDType): + df[config.time_column] = pd.to_datetime(df[config.time_column], utc=True) + elif isinstance(df[config.time_column].dtype, DatetimeTZDtype): + # Already tz-aware, convert to UTC + df[config.time_column] = df[config.time_column].dt.tz_convert("UTC") + else: + df[config.time_column] = df[config.time_column].dt.tz_localize("UTC") + else: + if isinstance(df[config.time_column].dtype, ObjectDType): + df[config.time_column] = pd.to_datetime(df[config.time_column], utc=False) + + +def _convert_spark_output_for_datetime(df: pd.DataFrame, config: DatetimeRanges) -> None: + """Convert DataFrame datetime columns after Spark output. + + Spark stores timestamps in UTC (as configured in SparkBackend). + This function converts them to match the expected timezone from the config. + """ + if config.time_column not in df.columns: + return + + # Ensure column is datetime (handle strings returned by Spark) + if not pd.api.types.is_datetime64_any_dtype(df[config.time_column]): + # Parse strings to datetime. Use utc=True to safely handle potential mixed offsets + # or explicit UTC strings. This results in a UTC-aware datetime series. + df[config.time_column] = pd.to_datetime(df[config.time_column], utc=True) + + if config.dtype == TimeDataType.TIMESTAMP_TZ: + # For tz-aware configs, ensure we have tz-aware timestamps. + # Spark timestamps are UTC. + if not isinstance(df[config.time_column].dtype, DatetimeTZDtype): + # tz-naive (default from Spark), localize to UTC + df[config.time_column] = df[config.time_column].dt.tz_localize("UTC") + + # Convert to target_tz if specified + target_tz = config.start.tzinfo + if target_tz is not None: + df[config.time_column] = df[config.time_column].dt.tz_convert(target_tz) + else: # TIMESTAMP_NTZ + # For tz-naive configs, ensure the column is tz-naive. + # If we just parsed with utc=True, we have UTC-aware datetimes. + # We need to strip the timezone using tz_convert(None) for tz-aware timestamps. + if isinstance(df[config.time_column].dtype, DatetimeTZDtype): + df[config.time_column] = df[config.time_column].dt.tz_convert(None) + + +def _write_to_duckdb( + backend: IbisBackend, + df: pd.DataFrame | pa.Table, + table_name: str, + if_exists: str, +) -> None: + """Write DataFrame or PyArrow Table to DuckDB.""" + match if_exists: + case "append": + backend.insert(table_name, df) + case "replace": + backend.drop_table(table_name, if_exists=True) + backend.create_table(table_name, df) + case "fail": + backend.create_table(table_name, df) + case _: + msg = f"Invalid if_exists value: {if_exists}" + raise InvalidOperation(msg) + + +def _write_to_sqlite( + backend: IbisBackend, + df: pd.DataFrame | pa.Table, + table_name: str, + configs: Sequence[TimeBaseModel], + if_exists: str, +) -> None: + """Write DataFrame to SQLite with datetime conversion.""" + _check_one_config_per_datetime_column(configs) + + # Convert PyArrow to pandas for datetime handling + if isinstance(df, pa.Table): + df = df.to_pandas() + + copied = False + for config in configs: + if isinstance(config, _DATETIME_RANGES): + df, copied = _convert_database_input_for_datetime(df, config, copied) + + match if_exists: + case "append": + backend.insert(table_name, df) + case "replace": + backend.drop_table(table_name, if_exists=True) + backend.create_table(table_name, df) + case "fail": + backend.create_table(table_name, df) + case _: + msg = f"Invalid if_exists value: {if_exists}" + raise InvalidOperation(msg) + + +def _write_to_pyspark( + backend: IbisBackend, + df: pd.DataFrame | pa.Table, + table_name: str, + configs: Sequence[TimeBaseModel], + if_exists: str, + scratch_dir: Path | None, +) -> None: + """Write DataFrame to Spark. + + Note: For Spark, we typically write to Parquet and create views. + Direct table writes are limited. + """ + # Convert PyArrow to pandas for datetime handling + if isinstance(df, pa.Table): + df = df.to_pandas() + + # datetime conversion is handled in backend.create_temp_view -> _prepare_data_for_spark + + match if_exists: + case "append": + msg = "INSERT INTO is not supported with Spark backend" + raise InvalidOperation(msg) + case "replace": + backend.drop_view(table_name, if_exists=True) + backend.create_temp_view(table_name, df) + case "fail": + backend.create_temp_view(table_name, df) + case _: + msg = f"Invalid if_exists value: {if_exists}" + raise InvalidOperation(msg) diff --git a/src/chronify/ibis/spark_backend.py b/src/chronify/ibis/spark_backend.py new file mode 100644 index 0000000..259dfdc --- /dev/null +++ b/src/chronify/ibis/spark_backend.py @@ -0,0 +1,241 @@ +"""Spark backend implementation.""" + +from pathlib import Path +from typing import Any, Literal, cast + +import ibis +import ibis.expr.types as ir +import pandas as pd +import pyarrow as pa + +from chronify.exceptions import InvalidOperation +from chronify.ibis.base import IbisBackend + + +class SparkBackend(IbisBackend): + """Spark backend implementation.""" + + def __init__( + self, + session: Any = None, + **kwargs: Any, + ) -> None: + """Create a Spark backend. + + Parameters + ---------- + session + Existing SparkSession, or None to create a new one + **kwargs + Additional arguments passed to ibis.pyspark.connect + """ + try: + import pyspark # noqa: F401 + except ImportError: + msg = ( + "PySpark is required for PySparkBackend. Install with: pip install chronify[spark]" + ) + raise ImportError(msg) + + if session is None: + from pyspark.sql import SparkSession + + session = SparkSession.builder.getOrCreate() + + connection = ibis.pyspark.connect(session, **kwargs) + super().__init__(connection) + self._session = session + + # Set timezone to UTC for consistency + session.conf.set("spark.sql.session.timeZone", "UTC") + # Use microsecond precision for Parquet timestamps + session.conf.set("spark.sql.parquet.outputTimestampType", "TIMESTAMP_MICROS") + + @property + def name(self) -> str: + return "spark" + + @property + def database(self) -> str | None: + return None + + def _prepare_data_for_spark(self, data: pd.DataFrame) -> pd.DataFrame: + """Prepare data for Spark by ensuring consistent timezone handling.""" + df = data.copy() + + # Convert timestamps to strings to avoid timezone confusion in Spark + # Spark will parse the strings according to session timezone (UTC) + # Naive "2020-01-01 12:00:00" -> 12:00:00 UTC + # Aware "2020-01-01 12:00:00-05:00" -> 17:00:00 UTC + for col in df.columns: + if pd.api.types.is_datetime64_any_dtype(df[col]): + df[col] = df[col].astype(str) + + return df + + def create_table( + self, + name: str, + data: pd.DataFrame | pa.Table | ir.Table, + overwrite: bool = False, + ) -> ir.Table: + mode: Literal["overwrite", "error"] = "overwrite" if overwrite else "error" + if isinstance(data, ir.Table): + sql = data.compile() + self._session.sql(f"CREATE TABLE {name} AS {sql}") + else: + df = self._prepare_data_for_spark(data) + spark_df = self._session.createDataFrame(df) + spark_df.write.saveAsTable(name, mode=mode) + + return self.table(name) + + def create_view(self, name: str, expr: ir.Table) -> None: + sql = expr.compile() + self._session.sql(f"CREATE VIEW {name} AS {sql}") + + def create_temp_view(self, name: str, data: ir.Table | pd.DataFrame | pa.Table) -> None: + if isinstance(data, ir.Table): + sql = data.compile() + self._session.sql(f"CREATE OR REPLACE TEMP VIEW {name} AS {sql}") + else: + df = self._prepare_data_for_spark(data) + spark_df = self._session.createDataFrame(df) + spark_df.createOrReplaceTempView(name) + + def drop_table(self, name: str, if_exists: bool = False) -> None: + if_exists_sql = "IF EXISTS " if if_exists else "" + try: + self._session.sql(f"DROP TABLE {if_exists_sql}{name}") + except Exception as e: + if "is a VIEW" in str(e): + self._session.sql(f"DROP VIEW {if_exists_sql}{name}") + else: + raise + + def drop_view(self, name: str, if_exists: bool = False) -> None: + if_exists_sql = "IF EXISTS " if if_exists else "" + self._session.sql(f"DROP VIEW {if_exists_sql}{name}") + + def list_tables(self) -> list[str]: + # Filter out internal Ibis memtables + return [ + t + for t in self._connection.list_tables() + if not t.startswith("ibis_pandas_memtable_") + and not t.startswith("ibis_pyarrow_memtable_") + and not t.startswith("_temp_") + and not t.startswith("_backing_") + ] + + def table(self, name: str) -> ir.Table: + return self._connection.table(name) + + def execute(self, expr: ir.Expr) -> pd.DataFrame: + # Spark's toPandas() has DST issues when converting timestamps. + # Work around by using Spark SQL to convert timestamps to strings, + # then convert back to pandas timestamps after toPandas(). + if isinstance(expr, ir.Table): + return self._execute_table_with_timestamp_fix(expr) + return cast(pd.DataFrame, expr.execute()) + + def _execute_table_with_timestamp_fix(self, expr: ir.Table) -> pd.DataFrame: + """Execute a table expression with workaround for Spark timestamp DST issues.""" + schema = expr.schema() + timestamp_cols = [name for name, dtype in schema.items() if dtype.is_timestamp()] + + if not timestamp_cols: + return cast(pd.DataFrame, expr.execute()) + + # Convert timestamp columns to strings in the query + sql = expr.compile() + select_exprs = [] + for col_name in schema.names: + if col_name in timestamp_cols: + # Cast to string to avoid DST conversion issues + select_exprs.append(f"CAST({col_name} AS STRING) AS {col_name}") + else: + select_exprs.append(col_name) + + query = f"SELECT {', '.join(select_exprs)} FROM ({sql})" + df: pd.DataFrame = self._session.sql(query).toPandas() + + # Convert string columns back to timestamps + for col_name in timestamp_cols: + df[col_name] = pd.to_datetime(df[col_name]) + + return df + + def insert(self, table_name: str, data: pd.DataFrame | pa.Table | ir.Table) -> None: + if isinstance(data, ir.Table): + sql = data.compile() + self._session.sql(f"INSERT INTO {table_name} {sql}") + else: + df = self._prepare_data_for_spark(data) + spark_df = self._session.createDataFrame(df) + spark_df.write.insertInto(table_name) + + def sql(self, query: str) -> ir.Table: + return self._connection.sql(query) + + def execute_sql_to_df(self, query: str) -> pd.DataFrame: + """Execute raw SQL and return results as a DataFrame. + + Uses string conversion for timestamp columns to avoid Spark's + DST conversion issues during toPandas(). + """ + spark_df = self._session.sql(query) + + # Check for timestamp columns + timestamp_cols = [ + field.name + for field in spark_df.schema.fields + if "timestamp" in field.dataType.simpleString().lower() + ] + + if not timestamp_cols: + return cast(pd.DataFrame, spark_df.toPandas()) + + # Convert timestamp columns to strings to avoid DST issues + from pyspark.sql.functions import col + + select_exprs = [] + for field in spark_df.schema.fields: + if field.name in timestamp_cols: + select_exprs.append(col(field.name).cast("string").alias(field.name)) + else: + select_exprs.append(col(field.name)) + + df: pd.DataFrame = spark_df.select(*select_exprs).toPandas() + + # Convert string columns back to timestamps + for col_name in timestamp_cols: + df[col_name] = pd.to_datetime(df[col_name]) + + return df + + def write_parquet( + self, + expr: ir.Table, + path: Path, + partition_by: list[str] | None = None, + overwrite: bool = False, + ) -> None: + if not overwrite and path.exists(): + msg = f"File already exists: {path}" + raise InvalidOperation(msg) + + sql = expr.compile() + if partition_by: + cols = ",".join(partition_by) + # Spark uses INSERT OVERWRITE DIRECTORY + query = ( + f"INSERT OVERWRITE DIRECTORY '{path}' USING parquet PARTITIONED BY ({cols}) {sql}" + ) + else: + query = f"INSERT OVERWRITE DIRECTORY '{path}' USING parquet {sql}" + self._session.sql(query) + + def create_view_from_parquet(self, name: str, path: Path) -> None: + query = f"CREATE VIEW {name} AS SELECT * FROM parquet.`{path}`" + self._session.sql(query) diff --git a/src/chronify/ibis/sqlite_backend.py b/src/chronify/ibis/sqlite_backend.py new file mode 100644 index 0000000..d84a81f --- /dev/null +++ b/src/chronify/ibis/sqlite_backend.py @@ -0,0 +1,156 @@ +"""SQLite backend implementation.""" + +from pathlib import Path +from typing import Any, Optional, cast + +import ibis +import ibis.expr.types as ir +import pandas as pd +import pyarrow as pa + +from chronify.exceptions import InvalidOperation +from chronify.ibis.base import IbisBackend + + +class SQLiteBackend(IbisBackend): + """SQLite backend implementation.""" + + def __init__( + self, + file_path: Optional[Path | str] = None, + **kwargs: Any, + ) -> None: + """Create a SQLite backend. + + Parameters + ---------- + file_path + Path to database file, or None for in-memory + **kwargs + Additional arguments passed to ibis.sqlite.connect + """ + self._file_path = str(file_path) if file_path else ":memory:" + connection = ibis.sqlite.connect( + self._file_path if self._file_path != ":memory:" else None, + **kwargs, + ) + super().__init__(connection) + + @property + def name(self) -> str: + return "sqlite" + + @property + def database(self) -> str | None: + return None if self._file_path == ":memory:" else self._file_path + + def create_table( + self, + name: str, + data: pd.DataFrame | pa.Table | ir.Table, + overwrite: bool = False, + ) -> ir.Table: + if overwrite and self.has_table(name): + self.drop_table(name, if_exists=True) + + if isinstance(data, ir.Table): + # SQLite loses datetime type info with CREATE TABLE AS SELECT + # Execute the query first and create from DataFrame to preserve types + df = data.execute() + self._connection.create_table(name, df) + else: + self._connection.create_table(name, data) + + return self.table(name) + + def create_view(self, name: str, expr: ir.Table) -> None: + sql = expr.compile() + self._connection.raw_sql(f"CREATE VIEW {name} AS {sql}") + + def create_temp_view(self, name: str, data: ir.Table | pd.DataFrame | pa.Table) -> None: + if isinstance(data, ir.Table): + sql = data.compile() + self._connection.raw_sql(f"CREATE TEMP VIEW {name} AS {sql}") + else: + # Create temp table first + self._connection.create_table(name, data, temp=True) + + def drop_table(self, name: str, if_exists: bool = False) -> None: + if_exists_sql = "IF EXISTS " if if_exists else "" + self._connection.raw_sql(f"DROP TABLE {if_exists_sql}{name}") + + def drop_view(self, name: str, if_exists: bool = False) -> None: + if_exists_sql = "IF EXISTS " if if_exists else "" + self._connection.raw_sql(f"DROP VIEW {if_exists_sql}{name}") + # Also drop the backing table if it exists (created by create_view_from_parquet) + backing_table = f"_parquet_backing_{name}" + self._connection.raw_sql(f"DROP TABLE IF EXISTS {backing_table}") + + def list_tables(self) -> list[str]: + # Filter out internal Ibis memtables and parquet backing tables + return [ + t + for t in self._connection.list_tables() + if not t.startswith("ibis_pandas_memtable_") and not t.startswith("_parquet_backing_") + ] + + def table(self, name: str) -> ir.Table: + return self._connection.table(name) + + def execute(self, expr: ir.Expr) -> pd.DataFrame: + return cast(pd.DataFrame, expr.execute()) + + def insert(self, table_name: str, data: pd.DataFrame | pa.Table | ir.Table) -> None: + if isinstance(data, ir.Table): + sql = data.compile() + self._connection.raw_sql(f"INSERT INTO {table_name} {sql}") + else: + # For pandas/pyarrow, use parameterized INSERT + if isinstance(data, pa.Table): + df = data.to_pandas() + else: + df = data + # Build parameterized INSERT statement + columns = ",".join(df.columns) + placeholders = ",".join(["?" for _ in df.columns]) + query = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})" + # Access SQLite's raw connection + cursor = self._connection.con.cursor() + cursor.executemany(query, df.values.tolist()) + self._connection.con.commit() + + def sql(self, query: str) -> ir.Table: + return self._connection.sql(query) + + def execute_sql_to_df(self, query: str) -> pd.DataFrame: + """Execute raw SQL and return results as a DataFrame.""" + cursor = self._connection.raw_sql(query) + columns = [desc[0] for desc in cursor.description] + data = cursor.fetchall() + return pd.DataFrame(data, columns=columns) + + def write_parquet( + self, + expr: ir.Table, + path: Path, + partition_by: Optional[list[str]] = None, + overwrite: bool = False, + ) -> None: + if partition_by: + msg = "SQLite does not support partitioned Parquet writes" + raise InvalidOperation(msg) + if not overwrite and path.exists(): + msg = f"File already exists: {path}" + raise InvalidOperation(msg) + + # Execute and write via pandas/pyarrow + df = self.execute(expr) + df.to_parquet(path) + + def create_view_from_parquet(self, name: str, path: Path) -> None: + # SQLite cannot directly read Parquet, so we load data into a backing table + # and create a view that selects from it + df = pd.read_parquet(path) + backing_table = f"_parquet_backing_{name}" + self._connection.create_table(backing_table, df) + self._connection.raw_sql(f"CREATE VIEW {name} AS SELECT * FROM {backing_table}") diff --git a/src/chronify/ibis/types.py b/src/chronify/ibis/types.py new file mode 100644 index 0000000..75fded6 --- /dev/null +++ b/src/chronify/ibis/types.py @@ -0,0 +1,206 @@ +"""Type conversion utilities for Ibis backends.""" + + +import duckdb +import ibis +import ibis.expr.datatypes as dt +import pandas as pd +import pyarrow as pa +from duckdb.typing import DuckDBPyType + +from chronify.exceptions import InvalidParameter + + +# Mapping from string type names to Ibis types +_COLUMN_TYPES = { + "bool": dt.Boolean(), + "datetime": dt.Timestamp(), + "float": dt.Float64(), + "int": dt.Int64(), + "bigint": dt.Int64(), + "str": dt.String(), +} + + +# Mapping from DuckDB type IDs to Ibis types +_DUCKDB_TO_IBIS_TYPES = { + duckdb.typing.BIGINT.id: dt.Int64(), # type: ignore + duckdb.typing.BOOLEAN.id: dt.Boolean(), # type: ignore + duckdb.typing.DOUBLE.id: dt.Float64(), # type: ignore + duckdb.typing.FLOAT.id: dt.Float32(), # type: ignore + duckdb.typing.INTEGER.id: dt.Int32(), # type: ignore + duckdb.typing.SMALLINT.id: dt.Int16(), # type: ignore + duckdb.typing.TINYINT.id: dt.Int8(), # type: ignore + duckdb.typing.VARCHAR.id: dt.String(), # type: ignore +} + + +def get_ibis_type_from_duckdb(duckdb_type: DuckDBPyType) -> dt.DataType: + """Return the Ibis type for a DuckDB type. + + Parameters + ---------- + duckdb_type + DuckDB type + + Returns + ------- + dt.DataType + Ibis data type + """ + match duckdb_type: + case duckdb.typing.TIMESTAMP_TZ: # type: ignore + return dt.Timestamp(timezone="UTC") + case ( + duckdb.typing.TIMESTAMP # type: ignore + | duckdb.typing.TIMESTAMP_MS # type: ignore + | duckdb.typing.TIMESTAMP_NS # type: ignore + | duckdb.typing.TIMESTAMP_S # type: ignore + ): + return dt.Timestamp() + case _: + ibis_type = _DUCKDB_TO_IBIS_TYPES.get(duckdb_type.id) + if ibis_type is None: + msg = f"There is no Ibis mapping for {duckdb_type=}" + raise InvalidParameter(msg) + return ibis_type + + +def get_ibis_type_from_string(type_name: str) -> dt.DataType: + """Return the Ibis type for a string type name. + + Parameters + ---------- + type_name + Type name (e.g., "int", "str", "datetime") + + Returns + ------- + dt.DataType + Ibis data type + """ + ibis_type = _COLUMN_TYPES.get(type_name) + if ibis_type is None: + msg = f"Unknown type name: {type_name}. Valid options: {list(_COLUMN_TYPES.keys())}" + raise InvalidParameter(msg) + return ibis_type + + +def get_ibis_types_from_dataframe( + df: pd.DataFrame | pa.Table, +) -> list[dt.DataType]: + """Return a list of Ibis types from a DataFrame or PyArrow table. + + Parameters + ---------- + df + Input DataFrame or PyArrow table + + Returns + ------- + list[dt.DataType] + List of Ibis types for each column + """ + # Use DuckDB to infer types (most robust approach) + if isinstance(df, pa.Table): + short_df = df.slice(0, 1) + else: + short_df = df.head(1) # noqa: F841 + + duckdb_types = duckdb.sql("SELECT * FROM short_df").dtypes + return [get_ibis_type_from_duckdb(t) for t in duckdb_types] + + +def get_ibis_schema_from_dataframe( + df: pd.DataFrame | pa.Table, +) -> ibis.Schema: + """Return an Ibis schema from a DataFrame or PyArrow table. + + Parameters + ---------- + df + Input DataFrame or PyArrow table + + Returns + ------- + ibis.Schema + Ibis schema + """ + if isinstance(df, pa.Table): + columns = df.column_names + else: + columns = list(df.columns) + + types = get_ibis_types_from_dataframe(df) + return ibis.schema(dict(zip(columns, types))) + + +def pyarrow_to_ibis_type(pa_type: pa.DataType) -> dt.DataType: + """Convert a PyArrow type to an Ibis type. + + Parameters + ---------- + pa_type + PyArrow data type + + Returns + ------- + dt.DataType + Ibis data type + """ + # Use Ibis's built-in conversion + return dt.from_pyarrow(pa_type) + + +def ibis_to_pyarrow_type(ibis_type: dt.DataType) -> pa.DataType: + """Convert an Ibis type to a PyArrow type. + + Parameters + ---------- + ibis_type + Ibis data type + + Returns + ------- + pa.DataType + PyArrow data type + """ + return ibis_type.to_pyarrow() + + +# Mapping from Ibis types to DuckDB types +_IBIS_TO_DUCKDB_TYPES: dict[type[dt.DataType], DuckDBPyType] = { + dt.Int8: duckdb.typing.TINYINT, # type: ignore + dt.Int16: duckdb.typing.SMALLINT, # type: ignore + dt.Int32: duckdb.typing.INTEGER, # type: ignore + dt.Int64: duckdb.typing.BIGINT, # type: ignore + dt.Float32: duckdb.typing.FLOAT, # type: ignore + dt.Float64: duckdb.typing.DOUBLE, # type: ignore + dt.Boolean: duckdb.typing.BOOLEAN, # type: ignore + dt.String: duckdb.typing.VARCHAR, # type: ignore + dt.Timestamp: duckdb.typing.TIMESTAMP, # type: ignore +} + + +def get_duckdb_type_from_ibis(ibis_type: dt.DataType) -> DuckDBPyType: + """Return the DuckDB type for an Ibis type. + + Parameters + ---------- + ibis_type + Ibis data type + + Returns + ------- + DuckDBPyType + DuckDB type + """ + # Check for timestamp with timezone + if isinstance(ibis_type, dt.Timestamp) and ibis_type.timezone is not None: + return duckdb.typing.TIMESTAMP_TZ # type: ignore + + duckdb_type = _IBIS_TO_DUCKDB_TYPES.get(type(ibis_type)) + if duckdb_type is None: + msg = f"There is no DuckDB mapping for {ibis_type=}" + raise InvalidParameter(msg) + return duckdb_type diff --git a/src/chronify/models.py b/src/chronify/models.py index b770200..545668f 100644 --- a/src/chronify/models.py +++ b/src/chronify/models.py @@ -1,16 +1,17 @@ import re from typing import Any, Optional -import duckdb.typing -import pandas as pd -from duckdb.typing import DuckDBPyType +import ibis.expr.datatypes as dt from pydantic import Field, field_validator, model_validator -from sqlalchemy import BigInteger, Boolean, DateTime, Double, Float, Integer, SmallInteger, String from typing_extensions import Annotated from chronify.base_models import ChronifyBaseModel from chronify.exceptions import InvalidParameter, InvalidValue from chronify.time_configs import TimeConfig +from chronify.ibis.types import ( + get_ibis_type_from_string, + _COLUMN_TYPES, +) REGEX_NAME_REQUIREMENT = re.compile(r"^\w+$") @@ -142,83 +143,8 @@ def list_columns(self) -> list[str]: return time_columns -# TODO: print example tables here. - -_COLUMN_TYPES = { - "bool": Boolean, - "datetime": DateTime, - "float": Double, - "int": Integer, - "bigint": BigInteger, - "str": String, -} - -_DB_TYPES = {x for x in _COLUMN_TYPES.values()} - -_DUCKDB_TYPES_TO_SQLALCHEMY_TYPES = { - duckdb.typing.BIGINT.id: BigInteger, # type: ignore - duckdb.typing.BOOLEAN.id: Boolean, # type: ignore - duckdb.typing.DOUBLE.id: Double, # type: ignore - duckdb.typing.FLOAT.id: Float, # type: ignore - duckdb.typing.INTEGER.id: Integer, # type: ignore - duckdb.typing.TINYINT.id: SmallInteger, # type: ignore - duckdb.typing.VARCHAR.id: String, # type: ignore - # Note: timestamp requires special handling because of timezone in sqlalchemy. -} - - -def get_sqlalchemy_type_from_duckdb(duckdb_type: DuckDBPyType) -> Any: - """Return the sqlalchemy type for a duckdb type.""" - match duckdb_type: - case duckdb.typing.TIMESTAMP_TZ: # type: ignore - sqlalchemy_type = DateTime(timezone=True) - case ( - duckdb.typing.TIMESTAMP # type: ignore - | duckdb.typing.TIMESTAMP_MS # type: ignore - | duckdb.typing.TIMESTAMP_NS # type: ignore - | duckdb.typing.TIMESTAMP_S # type: ignore - ): - sqlalchemy_type = DateTime(timezone=False) - case _: - cls = _DUCKDB_TYPES_TO_SQLALCHEMY_TYPES.get(duckdb_type.id) - if cls is None: - msg = f"There is no sqlalchemy mapping for {duckdb_type=}" - raise InvalidParameter(msg) - sqlalchemy_type = cls() - - return sqlalchemy_type - - -def get_duckdb_type_from_sqlalchemy(sqlalchemy_type: Any) -> DuckDBPyType: - """Return the duckdb type for a sqlalchemy type.""" - if isinstance(sqlalchemy_type, DateTime): - duckdb_type = ( - duckdb.typing.TIMESTAMP_TZ # type: ignore - if sqlalchemy_type.timezone - else duckdb.typing.TIMESTAMP # type: ignore - ) - elif isinstance(sqlalchemy_type, BigInteger): - duckdb_type = duckdb.typing.BIGINT # type: ignore - elif isinstance(sqlalchemy_type, Boolean): - duckdb_type = duckdb.typing.BOOLEAN # type: ignore - elif isinstance(sqlalchemy_type, Double): - duckdb_type = duckdb.typing.DOUBLE # type: ignore - elif isinstance(sqlalchemy_type, Integer): - duckdb_type = duckdb.typing.INTEGER # type: ignore - elif isinstance(sqlalchemy_type, String): - duckdb_type = duckdb.typing.VARCHAR # type: ignore - else: - msg = f"There is no duckdb mapping for {sqlalchemy_type=}" - raise InvalidParameter(msg) - - return duckdb_type # type: ignore - - -def get_duckdb_types_from_pandas(df: pd.DataFrame) -> list[DuckDBPyType]: - """Return a list of DuckDB types from a pandas dataframe.""" - # This seems least-prone to error, but is not exactly the most efficient. - short_df = df.head(1) # noqa: F841 - return duckdb.sql("select * from short_df").dtypes +# Set of valid Ibis dtype instances for validation +_IBIS_DTYPES = {type(v) for v in _COLUMN_TYPES.values()} class ColumnDType(ChronifyBaseModel): @@ -231,18 +157,24 @@ class ColumnDType(ChronifyBaseModel): @classmethod def fix_data_type(cls, data: dict[str, Any]) -> dict[str, Any]: dtype = data.get("dtype") - if dtype is None or any(map(lambda x: isinstance(dtype, x), _DB_TYPES)): + if dtype is None: + return data + + # If already an Ibis dtype instance, accept it + if isinstance(dtype, dt.DataType): return data if isinstance(dtype, str): - val = _COLUMN_TYPES.get(dtype) - if val is None: - options = sorted(_COLUMN_TYPES.keys()) + list(_DB_TYPES) + try: + data["dtype"] = get_ibis_type_from_string(dtype) + except InvalidParameter: + options = sorted(_COLUMN_TYPES.keys()) msg = f"{dtype=} must be one of {options}" raise InvalidValue(msg) - data["dtype"] = val() else: - msg = f"dtype is an unsupported type: {type(dtype)}. It must be a str or type." + msg = ( + f"dtype is an unsupported type: {type(dtype)}. It must be a str or Ibis DataType." + ) raise InvalidValue(msg) return data diff --git a/src/chronify/schema_manager.py b/src/chronify/schema_manager.py index 23828cc..bd14ab5 100644 --- a/src/chronify/schema_manager.py +++ b/src/chronify/schema_manager.py @@ -1,26 +1,11 @@ import json -from typing import Optional +import ibis.expr.types as ir from loguru import logger -from sqlalchemy import ( - Column, - Connection, - Engine, - MetaData, - String, - Table, - delete, - insert, - select, - text, -) - -from chronify.exceptions import ( - TableNotStored, -) -from chronify.models import ( - TableSchema, -) + +from chronify.exceptions import TableNotStored +from chronify.ibis.backend import IbisBackend +from chronify.models import TableSchema class SchemaManager: @@ -28,61 +13,59 @@ class SchemaManager: SCHEMAS_TABLE = "schemas" - def __init__(self, engine: Engine, metadata: MetaData): - self._engine = engine - self._metadata = metadata + def __init__(self, backend: IbisBackend): + self._backend = backend # Caching is not necessary if using SQLite, which provides very fast performance (~1 us) # for checking schemas in the **tiny** schemas table. # The same lookups in DuckDB are taking over 100 us. self._cache: dict[str, TableSchema] = {} - if self.SCHEMAS_TABLE in self._metadata.tables: - logger.info("Loaded existing database {}", self._engine.url.database) + if self.SCHEMAS_TABLE in self._backend.list_tables(): + logger.info("Loaded existing database {}", self._backend.database) self.rebuild_cache() else: - if self._engine.name == "hive": - # metadata.create_all doesn't work here. - with self._engine.begin() as conn: - conn.execute(text(f"DROP TABLE IF EXISTS {self.SCHEMAS_TABLE}")) - conn.execute( - text(f"CREATE TABLE {self.SCHEMAS_TABLE}(name STRING, schema STRING)") - ) - self._metadata.reflect(self._engine) - else: - table = Table( - self.SCHEMAS_TABLE, - self._metadata, - Column("name", String, nullable=False, unique=True), - Column("schema", String), # schema encoded as JSON - ) - self._metadata.create_all(self._engine, tables=[table]) - logger.info("Initialized new database: {}", self._engine.url.database) + self._create_schemas_table() + logger.info("Initialized new database: {}", self._backend.database) + + def _create_schemas_table(self) -> None: + """Create the schemas table.""" + # Use raw SQL to create table with proper types across all backends + if self._backend.name == "duckdb": + self._backend.execute_sql( + f"CREATE TABLE {self.SCHEMAS_TABLE}(name VARCHAR, schema VARCHAR)" + ) + elif self._backend.name == "spark": + # Spark 4.0+ requires VARCHAR to have a length; use STRING instead + self._backend.execute_sql( + f"CREATE TABLE {self.SCHEMAS_TABLE}(name STRING, schema STRING)" + ) + else: + # SQLite uses TEXT + self._backend.execute_sql(f"CREATE TABLE {self.SCHEMAS_TABLE}(name TEXT, schema TEXT)") - def _get_schema_table(self) -> Table: - return ( - self._metadata.tables[self.SCHEMAS_TABLE] - if self._engine.name == "hive" - else Table(self.SCHEMAS_TABLE, self._metadata) - ) + def _get_schema_table(self) -> ir.Table: + """Get the schemas table as an Ibis table.""" + return self._backend.table(self.SCHEMAS_TABLE) - def add_schema(self, conn: Connection, schema: TableSchema) -> Table: + def add_schema(self, schema: TableSchema) -> None: """Add the schema to the store.""" - table = self._get_schema_table() - stmt = insert(table).values(name=schema.name, schema=schema.model_dump_json()) - conn.execute(stmt) - # If there is a rollback after this addition to cached, things _should_ still be OK. + import pandas as pd + + df = pd.DataFrame({"name": [schema.name], "schema": [schema.model_dump_json()]}) + self._backend.insert(self.SCHEMAS_TABLE, df) + + # If there is a rollback after this addition to cache, things _should_ still be OK. # The table will be deleted and any attempted reads will fail with an error. # There will be a stale entry in cache, but it will be overwritten if the user ever # adds a new table with the same name. self._cache[schema.name] = schema logger.trace("Added schema for table {}", schema.name) - return table - def get_schema(self, name: str, conn: Optional[Connection] = None) -> TableSchema: + def get_schema(self, name: str) -> TableSchema: """Retrieve the schema for the table with name.""" schema = self._cache.get(name) if schema is None: - self.rebuild_cache(conn=conn) + self.rebuild_cache() schema = self._cache.get(name) if schema is None: @@ -91,43 +74,61 @@ def get_schema(self, name: str, conn: Optional[Connection] = None) -> TableSchem return self._cache[name] - def remove_schema(self, conn: Connection, name: str) -> None: + def remove_schema(self, name: str) -> None: """Remove the schema from the store.""" - table = self._get_schema_table() - if self._engine.name == "hive": - # Hive/Spark doesn't support delete, so we have to re-create the table without - # this one entry - stmt = select(table).where(table.c.name != name) - rows = conn.execute(stmt).fetchall() - conn.execute(text(f"DROP TABLE {self.SCHEMAS_TABLE}")) - conn.execute(text(f"CREATE TABLE {self.SCHEMAS_TABLE}(name STRING, schema STRING)")) - for row in rows: - params = {"name": row[0], "schema": row[1]} - conn.execute( - text(f"INSERT INTO {self.SCHEMAS_TABLE} VALUES(:name, :schema)"), - params, + if self._backend.name == "spark": + # Spark doesn't support DELETE, so we recreate the table + table = self._get_schema_table() + # Filter out the row to delete and get remaining rows + remaining = table.filter(table.name != name) + df = self._backend.execute(remaining) + + self._backend.execute_sql(f"DROP TABLE {self.SCHEMAS_TABLE}") + self._backend.execute_sql( + f"CREATE TABLE {self.SCHEMAS_TABLE}(name STRING, schema STRING)" + ) + for _, row in df.iterrows(): + escaped_name = _escape_sql_string(row["name"]) + escaped_schema = _escape_sql_string(row["schema"]) + self._backend.execute_sql( + f"INSERT INTO {self.SCHEMAS_TABLE} VALUES({escaped_name}, {escaped_schema})" ) else: - stmt2 = delete(table).where(table.c["name"] == name) - conn.execute(stmt2) + # DuckDB and SQLite support DELETE + escaped_name = _escape_sql_string(name) + self._backend.execute_sql( + f"DELETE FROM {self.SCHEMAS_TABLE} WHERE name = {escaped_name}" + ) - self._cache.pop(name) + self._cache.pop(name, None) - def rebuild_cache(self, conn: Optional[Connection] = None) -> None: + def rebuild_cache(self) -> None: """Rebuild the cache of schemas.""" self._cache.clear() - if conn is None: - with self._engine.connect() as conn: - self._rebuild_cache(conn) - else: - self._rebuild_cache(conn) - - def _rebuild_cache(self, conn: Connection) -> None: table = self._get_schema_table() - stmt = select(table) - res = conn.execute(stmt).fetchall() - for name, json_text in res: + df = self._backend.execute(table) + for _, row in df.iterrows(): + name = row["name"] + json_text = row["schema"] schema = TableSchema(**json.loads(json_text)) assert name == schema.name assert name not in self._cache self._cache[name] = schema + + +def _escape_sql_string(value: str) -> str: + """Escape a string value for safe inclusion in SQL statements. + + Parameters + ---------- + value + The string value to escape + + Returns + ------- + str + SQL-safe quoted string + """ + # Escape single quotes by doubling them + escaped = value.replace("'", "''") + return f"'{escaped}'" diff --git a/src/chronify/hive_functions.py b/src/chronify/spark_functions.py similarity index 64% rename from src/chronify/hive_functions.py rename to src/chronify/spark_functions.py index ff3ca07..e5efacc 100644 --- a/src/chronify/hive_functions.py +++ b/src/chronify/spark_functions.py @@ -2,18 +2,17 @@ from tempfile import NamedTemporaryFile from typing import Optional -from sqlalchemy import Engine, MetaData, text +from chronify.ibis.backend import IbisBackend def create_materialized_view( query: str, dst_table: str, - engine: Engine, - metadata: MetaData, + backend: IbisBackend, scratch_dir: Optional[Path] = None, ) -> None: """Create a materialized view with a Parquet file. This is a workaround for an undiagnosed - problem with timestamps and time zones with hive. + problem with timestamps and time zones with Spark/Hive. The Parquet file will be written to scratch_dir. Callers must ensure that the directory persists for the duration of the work. @@ -21,14 +20,13 @@ def create_materialized_view( with NamedTemporaryFile(dir=scratch_dir, suffix=".parquet") as f: f.close() output = Path(f.name) + write_query = f""" INSERT OVERWRITE DIRECTORY '{output}' USING parquet ({query}) """ - with engine.begin() as conn: - conn.execute(text(write_query)) - view_query = f"CREATE VIEW {dst_table} AS SELECT * FROM parquet.`{output}`" - conn.execute(text(view_query)) - metadata.reflect(engine, views=True) + backend.execute_sql(write_query) + view_query = f"CREATE VIEW {dst_table} AS SELECT * FROM parquet.`{output}`" + backend.execute_sql(view_query) diff --git a/src/chronify/sqlalchemy/__init__.py b/src/chronify/sqlalchemy/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/chronify/sqlalchemy/functions.py b/src/chronify/sqlalchemy/functions.py deleted file mode 100644 index 1e882f0..0000000 --- a/src/chronify/sqlalchemy/functions.py +++ /dev/null @@ -1,266 +0,0 @@ -"""This file provides functions to read and write the database as efficiently as possible. -The default behavior of sqlalchemy is to convert data into rows of tuples in Python, which -is very slow. This code attempts to bypass Python as much as possible through Arrow tables -in memory. -""" - -import atexit -from pathlib import Path -from tempfile import NamedTemporaryFile -from typing import Any, Literal, Optional, TypeAlias, Sequence -from collections import Counter - -import pandas as pd -from numpy.dtypes import DateTime64DType, ObjectDType -from pandas import DatetimeTZDtype -from chronify.time import TimeDataType -from sqlalchemy import Connection, Engine, Selectable, text - -from chronify.exceptions import InvalidOperation, InvalidParameter -from chronify.time_configs import DatetimeRangeBase, TimeBaseModel -from chronify.utils.path_utils import check_overwrite, delete_if_exists, to_path - -# Copied from Pandas/Polars -DbWriteMode: TypeAlias = Literal["replace", "append", "fail"] - - -def read_database( - query: Selectable | str, conn: Connection, config: TimeBaseModel, params: Any = None -) -> pd.DataFrame: - """Read a database query into a Pandas DataFrame.""" - match conn.engine.name: - case "duckdb": - if isinstance(query, str): - df = conn._dbapi_connection.driver_connection.sql(query, params=params).to_df() # type: ignore - else: - df = conn.execute(query).cursor.fetch_df() # type: ignore - case "sqlite": - df = pd.read_sql(query, conn, params=params) - if isinstance(config, DatetimeRangeBase): - _convert_database_output_for_datetime(df, config) - case "hive": - df = _read_from_hive(query, conn, config, params) - case _: - df = pd.read_sql(query, conn, params=params) - return df # type: ignore - - -def write_database( - df: pd.DataFrame, - conn: Connection, - table_name: str, - configs: Sequence[TimeBaseModel], - if_table_exists: DbWriteMode = "append", - scratch_dir: Path | None = None, -) -> None: - """Write a Pandas DataFrame to the database. - configs allows sqlite formatting for more than one datetime columns. - - Note: Writing persistent data with Hive as the backend is not supported. - This function will write the dataframe to a temporary Parquet file and then create - a view into that file. This is only to support ephemeral tables, such as for mapping tables. - """ - match conn.engine.name: - case "duckdb": - _write_to_duckdb(df, conn, table_name, if_table_exists) - case "sqlite": - _write_to_sqlite(df, conn, table_name, configs, if_table_exists) - case "hive": - _write_to_hive(df, conn, table_name, configs, if_table_exists, scratch_dir) - case _: - df.to_sql(table_name, conn, if_exists=if_table_exists, index=False) - - -def _check_one_config_per_datetime_column(configs: Sequence[TimeBaseModel]) -> None: - time_col_count = Counter( - [config.time_column for config in configs if isinstance(config, DatetimeRangeBase)] - ) - time_col_dup = {k: v for k, v in time_col_count.items() if v > 1} - if len(time_col_dup) > 0: - msg = f"More than one datetime config found for: {time_col_dup}" - raise InvalidParameter(msg) - - -def _convert_database_input_for_datetime( - df: pd.DataFrame, config: DatetimeRangeBase, copied: bool -) -> tuple[pd.DataFrame, bool]: - if config.dtype == TimeDataType.TIMESTAMP_NTZ: - return df, copied - - if copied: - df2 = df - else: - df2 = df.copy() - copied = True - if isinstance(df2[config.time_column].dtype, DatetimeTZDtype): - df2[config.time_column] = df2[config.time_column].dt.tz_convert("UTC") - else: - df2[config.time_column] = df2[config.time_column].dt.tz_localize("UTC") - - return df2, copied - - -def _convert_database_output_for_datetime(df: pd.DataFrame, config: DatetimeRangeBase) -> None: - if config.time_column in df.columns: - if config.dtype == TimeDataType.TIMESTAMP_TZ: - if isinstance(df[config.time_column].dtype, ObjectDType): - df[config.time_column] = pd.to_datetime(df[config.time_column], utc=True) - else: - df[config.time_column] = df[config.time_column].dt.tz_localize("UTC") - else: - if isinstance(df[config.time_column].dtype, ObjectDType): - df[config.time_column] = pd.to_datetime(df[config.time_column], utc=False) - - -def _write_to_duckdb( - df: pd.DataFrame, - conn: Connection, - table_name: str, - if_table_exists: DbWriteMode, -) -> None: - assert conn._dbapi_connection is not None - assert conn._dbapi_connection.driver_connection is not None - - match if_table_exists: - case "append": - query = f"INSERT INTO {table_name} SELECT * FROM df" - case "replace": - conn._dbapi_connection.driver_connection.sql(f"DROP TABLE IF EXISTS {table_name}") - query = f"CREATE TABLE {table_name} AS SELECT * FROM df" - case "fail": - query = f"CREATE TABLE {table_name} AS SELECT * FROM df" - case _: - msg = f"{if_table_exists=}" - raise InvalidOperation(msg) - - conn._dbapi_connection.driver_connection.sql(query) - - -def _write_to_hive( - df: pd.DataFrame, - conn: Connection, - table_name: str, - configs: Sequence[TimeBaseModel], - if_table_exists: DbWriteMode, - scratch_dir: Path | None, -) -> None: - df2 = df.copy() - for config in configs: - if isinstance(config, DatetimeRangeBase): - if isinstance(df2[config.time_column].dtype, DatetimeTZDtype): - # Spark doesn't like ns. That might change in the future. - # Pandas might offer a better way to change from ns to us in the future. - new_dtype = df2[config.time_column].dtype.name.replace( - "datetime64[ns", "datetime64[us" - ) - df2[config.time_column] = df2[config.time_column].astype(new_dtype) # type: ignore - elif isinstance(df2[config.time_column].dtype, DateTime64DType): - new_dtype = "datetime64[us]" - df2[config.time_column] = df2[config.time_column].astype(new_dtype) # type: ignore - else: - new_dtype = "datetime64[us]" - df2[config.time_column] = pd.to_datetime( - df2[config.time_column], utc=False, errors="raise" - ).astype(new_dtype) # type: ignore - - with NamedTemporaryFile(suffix=".parquet", dir=scratch_dir) as f: - f.close() - output = Path(f.name) - - df2.to_parquet(output) - atexit.register(lambda: delete_if_exists(output)) - select_stmt = f"SELECT * FROM parquet.`{output}`" - # TODO: CREATE TABLE causes DST fallback timestamps to get dropped - match if_table_exists: - case "append": - msg = "INSERT INTO is not supported with write_to_hive" - raise InvalidOperation(msg) - case "replace": - conn.execute(text(f"DROP VIEW IF EXISTS {table_name}")) - query = f"CREATE VIEW {table_name} AS {select_stmt}" - case "fail": - # Let the database fail the operation if the table already exists. - query = f"CREATE VIEW {table_name} AS {select_stmt}" - case _: - msg = f"{if_table_exists=}" - raise InvalidOperation(msg) - conn.execute(text(query)) - - -def _read_from_hive( - query: Selectable | str, conn: Connection, config: TimeBaseModel, params: Any = None -) -> pd.DataFrame: - df = pd.read_sql_query(query, conn, params=params) - if ( - isinstance(config, DatetimeRangeBase) - and config.time_column in df.columns - and not config.start_time_is_tz_naive() - ): - # This is tied to the fact that we set the Spark session to UTC. - # Otherwise, there is confusion with the computer's local time zone. - df[config.time_column] = df[config.time_column].dt.tz_localize("UTC") - return df - - -def _write_to_sqlite( - df: pd.DataFrame, - conn: Connection, - table_name: str, - configs: Sequence[TimeBaseModel], - if_table_exists: DbWriteMode, -) -> None: - _check_one_config_per_datetime_column(configs) - copied = False - for config in configs: - if isinstance(config, DatetimeRangeBase): - df, copied = _convert_database_input_for_datetime(df, config, copied) - df.to_sql(table_name, conn, if_exists=if_table_exists, index=False) - - -def create_view_from_parquet(conn: Connection, view_name: str, filename: Path) -> None: - """Create a view from a Parquet file.""" - if conn.engine.name == "duckdb": - str_path = f"{filename}/**/*.parquet" if filename.is_dir() else str(filename) - query = f"CREATE VIEW {view_name} AS SELECT * FROM read_parquet('{str_path}')" - elif conn.engine.name == "hive": - query = f"CREATE VIEW {view_name} AS SELECT * FROM parquet.`{filename}`" - else: - msg = f"create_view_from_parquet does not support engine={conn.engine.name}" - raise NotImplementedError(msg) - conn.execute(text(query)) - - -def write_query_to_parquet( - engine: Engine, - query: str, - output_file: Path, - overwrite: bool = False, - partition_columns: Optional[list[str]] = None, -) -> None: - """Write the query to a Parquet file.""" - output_file = to_path(output_file) - check_overwrite(output_file, overwrite) - match engine.name: - case "duckdb": - if partition_columns: - cols = ",".join(partition_columns) - query = ( - f"COPY ({query}) TO '{output_file}' (FORMAT PARQUET, PARTITION_BY ({cols}))" - ) - else: - query = f"COPY ({query}) TO '{output_file}' (FORMAT PARQUET)" - case "hive": - if not overwrite: - msg = "write_table_to_parquet with Hive requires overwrite=True" - raise InvalidOperation(msg) - # TODO: partition columns - if partition_columns: - msg = "write_table_to_parquet with Hive doesn't support partition_columns" - raise InvalidOperation(msg) - query = f"INSERT OVERWRITE DIRECTORY '{output_file}' USING parquet {query}" - case _: - msg = f"{engine.name=}" - raise NotImplementedError(msg) - - with engine.connect() as conn: - conn.execute(text(query)) diff --git a/src/chronify/store.py b/src/chronify/store.py index a6156ff..ab1fdc3 100644 --- a/src/chronify/store.py +++ b/src/chronify/store.py @@ -2,26 +2,14 @@ from pathlib import Path import shutil from typing import Any, Optional -from chronify.utils.sql import make_temp_view_name from datetime import tzinfo import duckdb import pandas as pd from duckdb import DuckDBPyRelation from loguru import logger -from sqlalchemy import ( - Column, - Connection, - Engine, - MetaData, - Selectable, - Table, - create_engine, - delete, - func, - select, - text, -) +import ibis.expr.types as ir +import pyarrow as pa import chronify.duckdb.functions as ddbf from chronify.exceptions import ( @@ -33,18 +21,17 @@ TableNotStored, ) from chronify.csv_io import read_csv +from chronify.ibis.backend import IbisBackend, make_backend +from chronify.ibis.functions import ( + create_view_from_parquet, + read_query, + write_parquet, + write_table, +) from chronify.models import ( CsvTableSchema, PivotedTableSchema, TableSchema, - get_duckdb_types_from_pandas, - get_sqlalchemy_type_from_duckdb, -) -from chronify.sqlalchemy.functions import ( - create_view_from_parquet, - read_database, - write_database, - write_query_to_parquet, ) from chronify.schema_manager import SchemaManager from chronify.time_configs import DatetimeRange, IndexTimeRangeBase, TimeBasedDataAdjustment @@ -53,7 +40,6 @@ from chronify.time_zone_converter import TimeZoneConverter, TimeZoneConverterByColumn from chronify.time_zone_localizer import TimeZoneLocalizer, TimeZoneLocalizerByColumn from chronify.utils.path_utils import check_overwrite, to_path -from chronify.utils.sqlalchemy_view import create_view class Store: @@ -61,8 +47,8 @@ class Store: def __init__( self, - engine: Optional[Engine] = None, - engine_name: Optional[str] = None, + backend: Optional[IbisBackend] = None, + backend_name: Optional[str] = None, file_path: Optional[Path | str] = None, **connect_kwargs: Any, ) -> None: @@ -70,10 +56,11 @@ def __init__( Parameters ---------- - engine - Optional, defaults to a engine connected to an in-memory DuckDB database. - engine_name - Optional, name of engine to use ('duckdb', 'sqlite'). Mutually exclusive with engine. + backend + Optional, an IbisBackend instance. Defaults to a DuckDB in-memory backend. + backend_name + Optional, name of backend to use ('duckdb', 'sqlite', 'spark'). + Mutually exclusive with backend. file_path Optional, use this file for the database. If the file does not exist, create a new database. If the file exists, load that existing database. @@ -81,109 +68,84 @@ def __init__( Examples -------- - >>> from sqlalchemy >>> store1 = Store() - >>> store2 = Store(engine=Engine("duckdb:///time_series.db")) - >>> store3 = Store(engine=Engine("sqlite:///time_series.db")) - >>> store4 = Store(engine_name="sqlite") + >>> store2 = Store(backend_name="duckdb", file_path="time_series.db") + >>> store3 = Store(backend_name="sqlite", file_path="time_series.db") + >>> store4 = Store(backend=make_backend("duckdb")) """ - self._metadata = MetaData() - if engine and engine_name: - msg = f"{engine=} and {engine_name=} cannot both be set" + if backend and backend_name: + msg = f"{backend=} and {backend_name=} cannot both be set" raise ConflictingInputsError(msg) - filename = ":memory:" if file_path is None else str(file_path) - if engine is None: - name = engine_name or "duckdb" - match name: - case "duckdb" | "sqlite": - engine_path = f"{name}:///{filename}" - case _: - msg = f"{engine_name=}" - raise NotImplementedError(msg) - self._engine = create_engine(engine_path, **connect_kwargs) + + if backend is None: + name = backend_name or "duckdb" + self._backend = make_backend(name, file_path=file_path, **connect_kwargs) else: - self._engine = engine + self._backend = backend - self._schema_mgr = SchemaManager(self._engine, self._metadata) - if self._engine.url.database != ":memory:": - self.update_metadata() + self._schema_mgr = SchemaManager(self._backend) @classmethod def create_in_memory_db( cls, - engine_name: str = "duckdb", + backend_name: str = "duckdb", **connect_kwargs: Any, ) -> "Store": """Create a Store with an in-memory database.""" - return Store(engine=create_engine(f"{engine_name}:///:memory:", **connect_kwargs)) + return Store(backend=make_backend(backend_name, **connect_kwargs)) @classmethod def create_file_db( cls, file_path: Path | str = "time_series.db", - engine_name: str = "duckdb", + backend_name: str = "duckdb", overwrite: bool = False, **connect_kwargs: Any, ) -> "Store": """Create a Store with a file-based database.""" path = to_path(file_path) check_overwrite(path, overwrite) - return Store(engine=create_engine(f"{engine_name}:///{path}", **connect_kwargs)) + return Store(backend=make_backend(backend_name, file_path=path, **connect_kwargs)) @classmethod - def create_new_hive_store( + def create_new_spark_store( cls, - url: str, + session: Any = None, drop_schema: bool = True, **connect_kwargs: Any, ) -> "Store": - """Create a new Store in a Hive database. + """Create a new Store with Spark backend. + Recommended usage is to create views from Parquet files. Ingesting data into tables from files or DataFrames is not supported. - This has been tested with Apache Spark running an Apache Thrift Server. - Parameters ---------- - url - Thrift server URL + session + Optional SparkSession. If None, a new one will be created. drop_schema If True, drop the schema table if it's already there. Examples -------- - >>> store = Store.create_new_hive_store("hive://localhost:10000/default") + >>> store = Store.create_new_spark_store() See also -------- create_view_from_parquet """ - # We don't currently expect to need to load an existing hive-based store, but it could - # be added. - if "hive://" not in url: - msg = f"Expected 'hive://' to be in url: {url}" - raise InvalidParameter(msg) - engine = create_engine(url, **connect_kwargs) - metadata = MetaData() - metadata.reflect(engine, views=True) - with engine.begin() as conn: - # Workaround for ambiguity of time zones in the read path. - conn.execute(text("SET TIME ZONE 'UTC'")) - # Workaround for the fact that Spark uses a non-standard format for timestamps - # in Parquet files. Pandas/DuckDB can't interpret them properly. - conn.execute(text("SET spark.sql.parquet.outputTimestampType=TIMESTAMP_MICROS")) - - if drop_schema: - if SchemaManager.SCHEMAS_TABLE in metadata.tables: - conn.execute(text(f"DROP TABLE {SchemaManager.SCHEMAS_TABLE}")) - - return cls(engine=engine) + backend = make_backend("spark", session=session, **connect_kwargs) + + if drop_schema and SchemaManager.SCHEMAS_TABLE in backend.list_tables(): + backend.drop_table(SchemaManager.SCHEMAS_TABLE) + + return cls(backend=backend) @classmethod def load_from_file( cls, file_path: Path | str, - engine_name: str = "duckdb", + backend_name: str = "duckdb", **connect_kwargs: Any, ) -> "Store": """Load an existing store from a database.""" @@ -191,83 +153,65 @@ def load_from_file( if not path.exists(): msg = str(path) raise FileNotFoundError(msg) - return Store(engine=create_engine(f"{engine_name}:///{path}", **connect_kwargs)) + return Store(backend=make_backend(backend_name, file_path=path, **connect_kwargs)) def dispose(self) -> None: - """Call self.engine.dispose() in order to dispose of the current connections.""" - self._engine.dispose() + """Clean up the backend connection.""" + self._backend.dispose() - def get_table(self, name: str) -> Table: - """Return the sqlalchemy Table object.""" + def get_table(self, name: str) -> ir.Table: + """Return the Ibis table reference.""" if not self.has_table(name): msg = f"{name=}" raise TableNotStored(msg) - return Table(name, self._metadata) + return self._backend.table(name) def has_table(self, name: str) -> bool: """Return True if the database has a table with the given name.""" - return name in self._metadata.tables + return self._backend.has_table(name) def list_tables(self) -> list[str]: """Return a list of user tables in the database.""" - return [x for x in self._metadata.tables if x != SchemaManager.SCHEMAS_TABLE] + return [x for x in self._backend.list_tables() if x != SchemaManager.SCHEMAS_TABLE] - def try_get_table(self, name: str) -> Table | None: - """Return the sqlalchemy Table object or None if it is not stored.""" + def try_get_table(self, name: str) -> ir.Table | None: + """Return the Ibis table object or None if it is not stored.""" if not self.has_table(name): return None - return Table(name, self._metadata) - - def update_metadata(self) -> None: - """Update the sqlalchemy metadata for table schema. Call this method if you add tables - in the sqlalchemy engine outside of this class or perform a rollback - in the same transaction in which chronify added tables. - """ - # Create a new object because sqlalchemy does not detect dropped tables in reflect. - metadata = MetaData() - metadata.reflect(self._engine, views=True) - logger.trace( - "Updated metadata, added: {}, dropped: {}", - sorted(set(metadata.tables).difference(self._metadata.tables)), - sorted(set(self._metadata.tables).difference(metadata.tables)), - ) - self._metadata = metadata - self._schema_mgr.rebuild_cache() + return self._backend.table(name) def backup(self, dst: Path | str, overwrite: bool = False) -> None: """Copy the database to a new location. Not yet supported for in-memory databases.""" - self._engine.dispose() + self._backend.dispose() path = to_path(dst) check_overwrite(path, overwrite) - match self._engine.name: - case "duckdb" | "sqlite": - if self._engine.url.database is None or self._engine.url.database == ":memory:": - msg = "backup is only supported with a database backed by a file" - raise InvalidOperation(msg) - src_file = Path(self._engine.url.database) - shutil.copyfile(src_file, path) - logger.info("Copied database to {}", path) - case _: - msg = self._engine.name - raise NotImplementedError(msg) - - @property - def engine(self) -> Engine: - """Return the sqlalchemy engine.""" - return self._engine + try: + match self._backend.name: + case "duckdb" | "sqlite": + if self._backend.database is None: + msg = "backup is only supported with a database backed by a file" + raise InvalidOperation(msg) + src_file = Path(self._backend.database) + shutil.copyfile(src_file, path) + logger.info("Copied database to {}", path) + case _: + msg = self._backend.name + raise NotImplementedError(msg) + finally: + self._backend.reconnect() @property - def metadata(self) -> MetaData: - """Return the sqlalchemy metadata.""" - return self._metadata + def backend(self) -> IbisBackend: + """Return the Ibis backend.""" + return self._backend @property def schema_manager(self) -> SchemaManager: """Return the store's schema manager.""" return self._schema_mgr - def check_timestamps(self, name: str, connection: Connection | None = None) -> None: + def check_timestamps(self, name: str) -> None: """Check the timestamps in the table. This is useful if you call a :meth:`ingest_table` many times with skip_time_checks=True @@ -285,25 +229,18 @@ def check_timestamps(self, name: str, connection: Connection | None = None) -> N """ table = self.get_table(name) schema = self._schema_mgr.get_schema(name) - if connection is None: - with self._engine.connect() as conn: - check_timestamps(conn, table, schema) - else: - check_timestamps(connection, table, schema) + check_timestamps(self._backend, table, schema) def create_view_from_parquet( self, path: Path, schema: TableSchema, bypass_checks: bool = False ) -> None: - """Load a table into the database.""" + """Load a table into the database from a Parquet file.""" self._create_view_from_parquet(path, schema) try: - with self._engine.connect() as conn: - table = self.get_table(schema.name) - if not bypass_checks: - check_timestamps(conn, table, schema) + table = self.get_table(schema.name) + if not bypass_checks: + check_timestamps(self._backend, table, schema) except InvalidTable: - # This doesn't use conn.rollback because we can't update the sqlalchemy metadata - # for this view inside the connection. self.drop_view(schema.name) raise @@ -340,18 +277,14 @@ def _create_view_from_parquet(self, path: Path | str, schema: TableSchema) -> No ... "table.parquet", ... ) """ - with self._engine.begin() as conn: - create_view_from_parquet(conn, schema.name, to_path(path)) - self._schema_mgr.add_schema(conn, schema) - - self.update_metadata() + create_view_from_parquet(self._backend, schema.name, to_path(path)) + self._schema_mgr.add_schema(schema) def ingest_from_csv( self, path: Path | str, src_schema: CsvTableSchema, dst_schema: TableSchema, - connection: Optional[Connection] = None, ) -> bool: """Ingest data from a CSV file. @@ -363,8 +296,6 @@ def ingest_from_csv( Defines the schema of the source file. dst_schema Defines the destination table in the database. - connection - Optional connection to reuse. Refer to :meth:`ingest_table` for notes. Returns ------- @@ -405,14 +336,13 @@ def ingest_from_csv( -------- ingest_from_csvs """ - return self.ingest_from_csvs((path,), src_schema, dst_schema, connection=connection) + return self.ingest_from_csvs((path,), src_schema, dst_schema) def ingest_from_csvs( self, paths: Iterable[Path | str], src_schema: CsvTableSchema, dst_schema: TableSchema, - connection: Optional[Connection] = None, ) -> bool: """Ingest data into the table specifed by schema. If the table does not exist, create it. This is faster than calling :meth:`ingest_from_csv` many times. @@ -422,14 +352,12 @@ def ingest_from_csvs( Parameters ---------- - path + paths Source data files src_schema Defines the schema of the source files. dst_schema Defines the destination table in the database. - conn - Optional connection to reuse. Refer to :meth:`ingest_table` for notes. Returns ------- @@ -446,46 +374,33 @@ def ingest_from_csvs( ingest_from_csv """ try: - if connection is None: - with self._engine.begin() as conn: - created_table = self._ingest_from_csvs(conn, paths, src_schema, dst_schema) - else: - created_table = self._ingest_from_csvs(connection, paths, src_schema, dst_schema) + created_table = self._ingest_from_csvs(paths, src_schema, dst_schema) except Exception: - # TODO: - # 1. The implicit rollback does not remove tables from our sqlalchemy metadata object. - # This means that the metadata object could be out-of-date if the user - # is self-managing the connection. - # 2. Python sqlite3 does not appear to support rollbacks with DDL statements. - # See discussion at https://bugs.python.org/issue10740. - self._handle_sqlite_error_case(dst_schema.name, connection) - if dst_schema.name in self._metadata.tables: - self._metadata.remove(Table(dst_schema.name, self._metadata)) + self._handle_error_case(dst_schema.name) raise return created_table def _ingest_from_csvs( self, - conn: Connection, paths: Iterable[Path | str], src_schema: CsvTableSchema, dst_schema: TableSchema, ) -> bool: created_table = False - if not paths: + paths_list = list(paths) + if not paths_list: return created_table - for path in paths: - if self._ingest_from_csv(conn, path, src_schema, dst_schema): + for path in paths_list: + if self._ingest_from_csv(path, src_schema, dst_schema): created_table = True - table = Table(dst_schema.name, self._metadata) - check_timestamps(conn, table, dst_schema) + table = self._backend.table(dst_schema.name) + check_timestamps(self._backend, table, dst_schema) return created_table def _ingest_from_csv( self, - conn: Connection, path: Path | str, src_schema: CsvTableSchema, dst_schema: TableSchema, @@ -494,36 +409,24 @@ def _ingest_from_csv( columns = set(src_schema.list_columns()) check_columns(rel.columns, columns) - # TODO if isinstance(src_schema.time_config, IndexTimeRangeBase): if isinstance(dst_schema.time_config, DatetimeRange): raise NotImplementedError - # timestamps = IndexTimeRangeGenerator(src_schema.time_config).list_timestamps() - # rel = ddbf.add_datetime_column( - # rel=rel, - # start=dst_schema.time_config.start, - # resolution=dst_schema.time_config.resolution, - # length=dst_schema.time_config.length, - # time_array_id_columns=src_schema.time_array_id_columns, - # time_column=dst_schema.time_config.time_column, - # timestamps=timestamps, - # ) else: cls_name = dst_schema.time_config.__class__.__name__ msg = f"{src_schema.time_config.__class__.__name__} cannot be converted to {cls_name}" raise NotImplementedError(msg) if src_schema.pivoted_dimension_name is not None: - return self._ingest_pivoted_table(conn, rel, src_schema, dst_schema) + return self._ingest_pivoted_table(rel, src_schema, dst_schema) - return self._ingest_table(conn, rel, dst_schema) + return self._ingest_table(rel, dst_schema) def ingest_pivoted_table( self, data: pd.DataFrame | DuckDBPyRelation, src_schema: PivotedTableSchema | CsvTableSchema, dst_schema: TableSchema, - connection: Optional[Connection] = None, ) -> bool: """Ingest pivoted data into the table specifed by schema. If the table does not exist, create it. Chronify will unpivot the data before ingesting it. @@ -536,8 +439,6 @@ def ingest_pivoted_table( Defines the schema of the input data. dst_schema Defines the destination table in the database. - conn - Optional connection to reuse. Refer to :meth:`ingest_table` for notes. Returns ------- @@ -588,14 +489,13 @@ def ingest_pivoted_table( -------- ingest_pivoted_tables """ - return self.ingest_pivoted_tables((data,), src_schema, dst_schema, connection=connection) + return self.ingest_pivoted_tables((data,), src_schema, dst_schema) def ingest_pivoted_tables( self, data: Iterable[pd.DataFrame | DuckDBPyRelation], src_schema: PivotedTableSchema | CsvTableSchema, dst_schema: TableSchema, - connection: Optional[Connection] = None, ) -> bool: """Ingest pivoted data into the table specifed by schema. @@ -612,8 +512,6 @@ def ingest_pivoted_tables( Defines the schema of all input tables. dst_schema Defines the destination table in the database. - conn - Optional connection to reuse. Refer to :meth:`ingest_table` for notes. Returns ------- @@ -625,38 +523,28 @@ def ingest_pivoted_tables( ingest_pivoted_table """ try: - if connection is None: - with self._engine.begin() as conn: - created_table = self._ingest_pivoted_tables(conn, data, src_schema, dst_schema) - else: - created_table = self._ingest_pivoted_tables( - connection, data, src_schema, dst_schema - ) + created_table = self._ingest_pivoted_tables(data, src_schema, dst_schema) except Exception: - self._handle_sqlite_error_case(dst_schema.name, connection) - if dst_schema.name in self._metadata.tables: - self._metadata.remove(Table(dst_schema.name, self._metadata)) + self._handle_error_case(dst_schema.name) raise return created_table def _ingest_pivoted_tables( self, - conn: Connection, data: Iterable[pd.DataFrame | DuckDBPyRelation], src_schema: PivotedTableSchema | CsvTableSchema, dst_schema: TableSchema, ) -> bool: created_table = False for table in data: - if self._ingest_pivoted_table(conn, table, src_schema, dst_schema): + if self._ingest_pivoted_table(table, src_schema, dst_schema): created_table = True - check_timestamps(conn, Table(dst_schema.name, self._metadata), dst_schema) + check_timestamps(self._backend, self._backend.table(dst_schema.name), dst_schema) return created_table def _ingest_pivoted_table( self, - conn: Connection, data: pd.DataFrame | DuckDBPyRelation, src_schema: PivotedTableSchema | CsvTableSchema, dst_schema: TableSchema, @@ -675,13 +563,12 @@ def _ingest_pivoted_table( src_schema.pivoted_dimension_name, dst_schema.value_column, ) - return self._ingest_table(conn, rel2, dst_schema) + return self._ingest_table(rel2, dst_schema) def ingest_table( self, - data: pd.DataFrame | DuckDBPyRelation, + data: pd.DataFrame | DuckDBPyRelation | ir.Table | pa.Table, schema: TableSchema, - connection: Optional[Connection] = None, **kwargs: Any, ) -> bool: """Ingest data into the table specifed by schema. If the table does not exist, @@ -693,13 +580,6 @@ def ingest_table( Input data to ingest into the database. schema Defines the destination table in the database. - connection - Optional connection to reuse. If adding many tables at once, it is significantly - faster to use one connection. Refer to :meth:`ingest_tables` for built-in support. - If connection is not set, chronify will commit the database changes - or perform a rollback on error. If it is set, the caller must perform those actions. - If you peform a rollback, you must call :meth:`rebuild_schema_cache` because the - Store will cache all table names in memory. Returns ------- @@ -743,13 +623,12 @@ def ingest_table( -------- ingest_tables """ - return self.ingest_tables((data,), schema, connection=connection, **kwargs) + return self.ingest_tables((data,), schema, **kwargs) def ingest_tables( self, - data: Iterable[pd.DataFrame | DuckDBPyRelation], + data: Iterable[pd.DataFrame | DuckDBPyRelation | ir.Table | pa.Table], schema: TableSchema, - connection: Optional[Connection] = None, **kwargs: Any, ) -> bool: """Ingest multiple input tables to the same database table. @@ -763,8 +642,6 @@ def ingest_tables( Input tables to ingest into one database table. schema Defines the destination table. - conn - Optional connection to reuse. Refer to :meth:`ingest_table` for notes. Returns ------- @@ -781,68 +658,72 @@ def ingest_tables( ingest_table """ created_table = False - if not data: + data_list = list(data) + if not data_list: return created_table + # Track if table existed before ingestion to avoid dropping existing data on error + table_existed_before = self.has_table(schema.name) + try: - if connection is None: - with self._engine.begin() as conn: - created_table = self._ingest_tables(conn, data, schema, **kwargs) - else: - created_table = self._ingest_tables(connection, data, schema, **kwargs) + created_table = self._ingest_tables(data_list, schema, **kwargs) except Exception: - self._handle_sqlite_error_case(schema.name, connection) - if schema.name in self._metadata.tables: - self._metadata.remove(Table(schema.name, self._metadata)) + # Only drop the table if we created it; don't destroy existing data + if not table_existed_before: + self._handle_error_case(schema.name) raise return created_table def _ingest_tables( self, - conn: Connection, - data: Iterable[pd.DataFrame | DuckDBPyRelation], + data: Iterable[pd.DataFrame | DuckDBPyRelation | ir.Table | pa.Table], schema: TableSchema, skip_time_checks: bool = False, ) -> bool: created_table = False for table in data: - if self._ingest_table(conn, table, schema): + if self._ingest_table(table, schema): created_table = True if not skip_time_checks: - check_timestamps(conn, Table(schema.name, self._metadata), schema) + check_timestamps(self._backend, self._backend.table(schema.name), schema) return created_table def _ingest_table( self, - conn: Connection, - data: pd.DataFrame | DuckDBPyRelation, + data: pd.DataFrame | DuckDBPyRelation | ir.Table | pa.Table, schema: TableSchema, ) -> bool: - if self._engine.name == "hive": - msg = "Data ingestion through Hive is not supported" + if self._backend.name == "spark": + msg = "Data ingestion through Spark is not supported" raise NotImplementedError(msg) - df = data.to_df() if isinstance(data, DuckDBPyRelation) else data - check_columns(df.columns, schema.list_columns()) - - table = self.try_get_table(schema.name) - if table is None: - duckdb_types = get_duckdb_types_from_pandas(df) - dtypes = [get_sqlalchemy_type_from_duckdb(x) for x in duckdb_types] - table = Table( - schema.name, - self._metadata, - *[Column(x, y) for x, y in zip(df.columns, dtypes)], - ) - self._metadata.create_all(conn) + + # Convert to DataFrame/PyArrow for consistent handling + if isinstance(data, ir.Table): + df = data.to_pyarrow() + elif isinstance(data, DuckDBPyRelation): + df = data.to_df() + else: + df = data + + if isinstance(df, pa.Table): + columns = df.column_names + else: + columns = df.columns + + check_columns(columns, schema.list_columns()) + + table_exists = self.has_table(schema.name) + if not table_exists: + # Use write_table to ensure proper datetime handling for all backends + write_table(self._backend, df, schema.name, [schema.time_config], if_exists="fail") created_table = True else: + write_table(self._backend, df, schema.name, [schema.time_config], if_exists="append") created_table = False - write_database(df, conn, schema.name, [schema.time_config]) - if created_table: - self._schema_mgr.add_schema(conn, schema) + self._schema_mgr.add_schema(schema) return created_table @@ -932,8 +813,7 @@ def map_table_time_config( src_schema = self._schema_mgr.get_schema(src_name) map_time( - self._engine, - self._metadata, + self._backend, src_schema, dst_schema, data_adjustment=data_adjustment, @@ -942,8 +822,7 @@ def map_table_time_config( output_file=output_file, check_mapped_timestamps=check_mapped_timestamps, ) - with self._engine.begin() as conn: - self._schema_mgr.add_schema(conn, dst_schema) + self._schema_mgr.add_schema(dst_schema) def convert_time_zone( self, @@ -1012,7 +891,7 @@ def convert_time_zone( """ src_schema = self._schema_mgr.get_schema(src_name) - tzc = TimeZoneConverter(self._engine, self._metadata, src_schema, time_zone) + tzc = TimeZoneConverter(self._backend, src_schema, time_zone) dst_schema = tzc.generate_to_schema() if self.has_table(dst_schema.name): @@ -1025,8 +904,7 @@ def convert_time_zone( check_mapped_timestamps=check_mapped_timestamps, ) - with self._engine.begin() as conn: - self._schema_mgr.add_schema(conn, dst_schema) + self._schema_mgr.add_schema(dst_schema) return dst_schema @@ -1106,7 +984,7 @@ def convert_time_zone_by_column( src_schema = self._schema_mgr.get_schema(src_name) tzc = TimeZoneConverterByColumn( - self._engine, self._metadata, src_schema, time_zone_column, wrap_time_allowed + self._backend, src_schema, time_zone_column, wrap_time_allowed ) dst_schema = tzc.generate_to_schema() @@ -1120,8 +998,7 @@ def convert_time_zone_by_column( check_mapped_timestamps=check_mapped_timestamps, ) - with self._engine.begin() as conn: - self._schema_mgr.add_schema(conn, dst_schema) + self._schema_mgr.add_schema(dst_schema) return dst_schema @@ -1192,7 +1069,7 @@ def localize_time_zone( """ src_schema = self._schema_mgr.get_schema(src_name) - tzl = TimeZoneLocalizer(self._engine, self._metadata, src_schema, time_zone) + tzl = TimeZoneLocalizer(self._backend, src_schema, time_zone) dst_schema = tzl.generate_to_schema() if self.has_table(dst_schema.name): @@ -1205,8 +1082,7 @@ def localize_time_zone( check_mapped_timestamps=check_mapped_timestamps, ) - with self._engine.begin() as conn: - self._schema_mgr.add_schema(conn, dst_schema) + self._schema_mgr.add_schema(dst_schema) return dst_schema @@ -1280,7 +1156,7 @@ def localize_time_zone_by_column( """ src_schema = self._schema_mgr.get_schema(src_name) - tzl = TimeZoneLocalizerByColumn(self._engine, self._metadata, src_schema, time_zone_column) + tzl = TimeZoneLocalizerByColumn(self._backend, src_schema, time_zone_column) dst_schema = tzl.generate_to_schema() if self.has_table(dst_schema.name): @@ -1293,17 +1169,14 @@ def localize_time_zone_by_column( check_mapped_timestamps=check_mapped_timestamps, ) - with self._engine.begin() as conn: - self._schema_mgr.add_schema(conn, dst_schema) + self._schema_mgr.add_schema(dst_schema) return dst_schema def read_query( self, name: str, - query: Selectable | str, - params: Any = None, - connection: Optional[Connection] = None, + query: ir.Table | str, ) -> pd.DataFrame: """Return the query result as a pandas DataFrame. @@ -1312,37 +1185,30 @@ def read_query( name Table or view name query - SQL query expressed as a string or salqlchemy Selectable - params - Parameters for SQL query if expressed as a string + SQL query as string or Ibis table expression Examples -------- - >>> df = store.read_query("SELECT * FROM devices") - >>> df = store.read_query("SELECT * FROM devices WHERE id = ?", params=(3,)) - - >>> from sqlalchemy import select - >>> table = store.schemas.get_table("devices") - >>> df = store.read_query(select(table).where(table.c.id == 3) + >>> df = store.read_query("devices", "SELECT * FROM devices") + >>> df = store.read_query("devices", store.get_table("devices").filter(...)) """ - schema = self._schema_mgr.get_schema(name, conn=connection) - if connection is None: - with self._engine.begin() as conn: - return read_database(query, conn, schema.time_config, params=params) + schema = self._schema_mgr.get_schema(name) + if isinstance(query, str): + expr = self._backend.sql(query) else: - return read_database(query, connection, schema.time_config, params=params) + expr = query + return read_query(self._backend, expr, schema.time_config) - def read_table(self, name: str, connection: Optional[Connection] = None) -> pd.DataFrame: + def read_table(self, name: str) -> pd.DataFrame: """Return the table as a pandas DataFrame.""" table = self.get_table(name) - stmt = select(table) - return self.read_query(name, stmt, connection=connection) + return self.read_query(name, table) def read_raw_query( - self, query: str, params: Any = None, connection: Optional[Connection] = None + self, + query: str, ) -> pd.DataFrame: - """Execute a query directly on the backend database connection, bypassing sqlalchemy, and - return the results as a DataFrame. + """Execute a query directly on the backend and return results as DataFrame. Note: Unlike :meth:`read_query`, no conversion of timestamps is performed. Timestamps will be in the format of the underlying database. SQLite backends will return @@ -1352,68 +1218,35 @@ def read_raw_query( ---------- query SQL query to execute - params - Optional parameters for SQL query - conn - Optional sqlalchemy connection returned by `Store.engine.connect()`. This can - improve performance when performing many reads. If used for database modifications, - it is the caller's responsibility to perform a commit and ensure that the connection - is closed correctly. Use of sqlalchemy's context manager is recommended. Examples -------- >>> store = Store() - >>> query1 = "SELECT * from my_table WHERE column = ?" - >>> params1 = ("value1",) - >>> query2 = "SELECT * from my_table WHERE column = ?'" - >>> params2 = ("value2",) - - >>> df = store.read_raw_query(query1, params=params1) - - >>> with store.engine.connect() as conn: - ... df1 = store.read_raw_query(query1, params=params1, connection=conn) - ... df2 = store.read_raw_query(query2, params=params2, connection=conn) + >>> df = store.read_raw_query("SELECT * from my_table WHERE column = 'value'") """ - if connection is None: - with self._engine.connect() as conn: - return self._read_raw_query(query, params, conn) - else: - return self._read_raw_query(query, params, connection) - - def _read_raw_query(self, query: str, params: Any, conn: Connection) -> pd.DataFrame: - assert conn._dbapi_connection is not None - assert conn._dbapi_connection.driver_connection is not None - match self._engine.name: - case "duckdb": - df = conn._dbapi_connection.driver_connection.sql(query, params=params).to_df() - assert isinstance(df, pd.DataFrame) - return df - case "sqlite": - return pd.read_sql(query, conn._dbapi_connection.driver_connection, params=params) - case _: - msg = self._engine.name - raise NotImplementedError(msg) + expr = self._backend.sql(query) + return self._backend.execute(expr) def write_query_to_parquet( self, - stmt: Selectable, + query: ir.Table | str, file_path: Path | str, overwrite: bool = False, partition_columns: Optional[list[str]] = None, ) -> None: """Write the result of a query to a Parquet file.""" - # We could add a separate path where the query is a string and skip the intermediate - # view if we passed parameters through the call stack. - view_name = make_temp_view_name() - create_view(view_name, stmt, self._engine, self._metadata) - try: - self.write_table_to_parquet( - view_name, file_path, overwrite=overwrite, partition_columns=partition_columns - ) - finally: - with self._engine.connect() as conn: - conn.execute(text(f"DROP VIEW {view_name}")) - self._metadata.remove(Table(view_name, self._metadata)) + if isinstance(query, str): + expr = self._backend.sql(query) + else: + expr = query + + write_parquet( + self._backend, + expr, + to_path(file_path), + overwrite=overwrite, + partition_columns=partition_columns, + ) def write_table_to_parquet( self, @@ -1427,9 +1260,10 @@ def write_table_to_parquet( msg = f"table {name=} is not stored" raise TableNotStored(msg) - write_query_to_parquet( - self._engine, - f"SELECT * FROM {name}", + table = self.get_table(name) + write_parquet( + self._backend, + table, to_path(file_path), overwrite=overwrite, partition_columns=partition_columns, @@ -1440,7 +1274,6 @@ def delete_rows( self, name: str, time_array_id_values: dict[str, Any], - connection: Optional[Connection] = None, ) -> int: """Delete all rows matching the time_array_id_values. @@ -1450,8 +1283,6 @@ def delete_rows( Name of table time_array_id_values Values for the time_array_id_values. Keys must match the columns in the schema. - connnection - Optional connection to the database. Refer :meth:`ingest_table` for notes. Returns ------- @@ -1462,16 +1293,12 @@ def delete_rows( -------- >>> store.delete_rows("devices", {"id": 47}) """ - # TODO: consider supporting a user-defined query. Would need to check consistency - # afterwards. - # The current approach doesn't need to check because only one single complete time - # array can be deleted on each call. table = self.get_table(name) if not time_array_id_values: msg = "time_array_id_values cannot be empty" raise InvalidParameter(msg) - schema = self._schema_mgr.get_schema(name, conn=connection) + schema = self._schema_mgr.get_schema(name) if sorted(time_array_id_values.keys()) != sorted(schema.time_array_id_columns): msg = ( "The keys of time_array_id_values must match the schema columns. " @@ -1480,27 +1307,30 @@ def delete_rows( ) raise InvalidParameter(msg) - assert time_array_id_values - stmt = delete(table) - - # duckdb does not offer a way to retrieve the number of deleted rows, so we must - # compute it manually. - # Deletions are not common. We are trading accuracy for peformance. - count_stmt = select(func.count()).select_from(table) - + # Build filter condition + filter_expr = None for column, value in time_array_id_values.items(): - stmt = stmt.where(table.c[column] == value) - count_stmt = count_stmt.where(table.c[column] == value) + condition = table[column] == value + filter_expr = condition if filter_expr is None else (filter_expr & condition) - if connection is None: - with self._engine.begin() as conn: - num_deleted = self._run_delete(conn, stmt, count_stmt) + # Count rows before deletion + count_result = self._backend.execute(table.filter(filter_expr).count()) + if isinstance(count_result, pd.DataFrame): + count_before = int(count_result.iloc[0, 0]) # type: ignore[arg-type] else: - num_deleted = self._run_delete(connection, stmt, count_stmt) - # Let the caller commit or rollback when ready. + count_before = int(count_result) + + # Build and execute DELETE statement + where_clauses = [ + f"{col} = {_escape_sql_value(val)}" for col, val in time_array_id_values.items() + ] + where_str = " AND ".join(where_clauses) + self._backend.execute_sql(f"DELETE FROM {name} WHERE {where_str}") + + num_deleted = count_before if num_deleted < 1: - msg = f"Failed to delete rows: {stmt=} {num_deleted=}" + msg = f"Failed to delete rows: {time_array_id_values=} {num_deleted=}" raise InvalidParameter(msg) logger.info( @@ -1509,88 +1339,98 @@ def delete_rows( time_array_id_values, ) - stmt2 = select(table).limit(1) - is_empty = False - if connection is None: - with self._engine.connect() as conn: - res = conn.execute(stmt2).fetchall() + # Check if table is empty + remaining_result = self._backend.execute(table.count()) + if isinstance(remaining_result, pd.DataFrame): + remaining_count = int(remaining_result.iloc[0, 0]) # type: ignore[arg-type] else: - res = connection.execute(stmt2).fetchall() - - if not res: - is_empty = True + remaining_count = int(remaining_result) + is_empty = remaining_count == 0 if is_empty: logger.info("Delete empty table {}", name) - self.drop_table(name, connection=connection) + self.drop_table(name) return num_deleted - def _run_delete(self, conn: Connection, stmt: Any, count_stmt: Any) -> int: - count1: int | None = None - if self._engine.name == "duckdb": - res1 = conn.execute(count_stmt).fetchone() - assert res1 is not None - count1 = res1[0] - res = conn.execute(stmt) - if self._engine.name == "duckdb": - res2 = conn.execute(count_stmt).fetchone() - assert res2 is not None - count2 = res2[0] - assert count1 is not None - num_deleted = count1 - count2 - else: - num_deleted = res.rowcount - return num_deleted # type: ignore - def drop_table( self, name: str, - connection: Optional[Connection] = None, if_exists: bool = False, ) -> None: """Drop a table from the database.""" - self._drop_table_or_view(name, "TABLE", connection, if_exists) + if not if_exists and not self.has_table(name): + msg = f"{name=}" + raise TableNotStored(msg) - def create_view(self, schema: TableSchema, stmt: Selectable) -> None: - """Create a view in the database.""" - create_view(schema.name, stmt, self._engine, self._metadata) - with self._engine.begin() as conn: - self._schema_mgr.add_schema(conn, schema) + self._backend.drop_table(name, if_exists=if_exists) + if name in self._schema_mgr._cache: + self._schema_mgr.remove_schema(name) + logger.info("Dropped table {}", name) - def drop_view( + def create_view( self, - name: str, - connection: Optional[Connection] = None, - if_exists: bool = False, + schema: TableSchema, + table: ir.Table, + bypass_checks: bool = False, ) -> None: - """Drop a view from the database.""" - self._drop_table_or_view(name, "VIEW", connection, if_exists) + """Register an Ibis table as a view with time validation. - def _drop_table_or_view( + Creates a view and persists the schema. Both the view and schema + will survive process restarts (for backends that support persistent storage). + + When using Spark, enabling the Hive metastore is required to persist restarts. + + Parameters + ---------- + schema + Defines the schema of the view. The name in the schema will be used as + the view name. + table + Ibis table expression to register. This can be a table loaded from a + Parquet file, a query result, or any other Ibis table expression. + bypass_checks + If True, skip time validation checks. Defaults to False. + + Raises + ------ + InvalidTable + Raised if time validation fails and bypass_checks is False. + """ + self._backend.create_view(schema.name, table) + self._schema_mgr.add_schema(schema) + + if not bypass_checks: + try: + registered_table = self.get_table(schema.name) + check_timestamps(self._backend, registered_table, schema) + except InvalidTable: + self._backend.drop_view(schema.name, if_exists=True) + self._backend.drop_table(f"_backing_{schema.name}", if_exists=True) + self._schema_mgr.remove_schema(schema.name) + raise + + logger.info("Registered permanent view {} with schema", schema.name) + + def drop_view( self, name: str, - table_type: str, - connection: Optional[Connection], - if_exists: bool, + if_exists: bool = False, ) -> None: - table = self.get_table(name) - if_exists_str = " IF EXISTS" if if_exists else "" - if connection is None: - with self._engine.begin() as conn: - conn.execute(text(f"DROP {table_type} {if_exists_str} {name}")) - self._schema_mgr.remove_schema(conn, name) - else: - connection.execute(text(f"DROP {table_type} {if_exists_str} {name}")) - self._schema_mgr.remove_schema(connection, name) + """Drop a view from the database.""" + if not if_exists and not self.has_table(name): + msg = f"{name=}" + raise TableNotStored(msg) - self._metadata.remove(table) - logger.info("Dropped {} {}", table_type.lower(), name) + self._backend.drop_view(name, if_exists=if_exists) + if name in self._schema_mgr._cache: + self._schema_mgr.remove_schema(name) + logger.info("Dropped view {}", name) - def _handle_sqlite_error_case(self, name: str, connection: Optional[Connection]) -> None: - if connection is None and self._engine.name == "sqlite": - with self._engine.begin() as conn: - conn.execute(text(f"DROP TABLE IF EXISTS {name}")) + def _handle_error_case(self, name: str) -> None: + """Clean up after an error during ingestion.""" + # Ibis backends don't support transactions for DDL, so manually drop the table + self._backend.drop_table(name, if_exists=True) def check_columns( @@ -1610,3 +1450,31 @@ def check_columns( cols = " ".join(sorted(diff)) msg = f"These columns are defined in the schema but not present in the table: {cols}" raise InvalidTable(msg) + + +def _escape_sql_value(value: Any) -> str: + """Escape a value for safe inclusion in SQL statements. + + Parameters + ---------- + value + The value to escape + + Returns + ------- + str + SQL-safe string representation of the value + """ + if value is None: + return "NULL" + if isinstance(value, bool): + return "TRUE" if value else "FALSE" + if isinstance(value, (int, float)): + return str(value) + if isinstance(value, str): + # Escape single quotes by doubling them + escaped = value.replace("'", "''") + return f"'{escaped}'" + # For other types, convert to string and escape + escaped = str(value).replace("'", "''") + return f"'{escaped}'" diff --git a/src/chronify/time.py b/src/chronify/time.py index 26d273f..8f3f70c 100644 --- a/src/chronify/time.py +++ b/src/chronify/time.py @@ -30,6 +30,10 @@ class TimeDataType(StrEnum): STRING = "string" TIMESTAMP_IN_PARTS = "timestamp_in_parts" + def is_timestamp(self) -> bool: + """Return True if the type is a timestamp type.""" + return self in (TimeDataType.TIMESTAMP_TZ, TimeDataType.TIMESTAMP_NTZ) + class RepresentativePeriodFormat(StrEnum): """Defines the supported formats for representative period data.""" diff --git a/src/chronify/time_configs.py b/src/chronify/time_configs.py index 2cb9f49..cae40a8 100644 --- a/src/chronify/time_configs.py +++ b/src/chronify/time_configs.py @@ -126,6 +126,7 @@ def check_dtype_start_time_consistency( f"\n{info.data['start']=}, {dtype=}" ) raise InvalidValue(msg) + assert dtype is not None # All valid cases are handled above return dtype diff --git a/src/chronify/time_series_checker.py b/src/chronify/time_series_checker.py index afaef59..b3a8bde 100644 --- a/src/chronify/time_series_checker.py +++ b/src/chronify/time_series_checker.py @@ -1,13 +1,14 @@ -from sqlalchemy import Connection, Table, select, text from typing import Optional from datetime import datetime, tzinfo +import ibis.expr.types as ir import pandas as pd from chronify.exceptions import InvalidTable +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import read_query from chronify.models import TableSchema from chronify.time_configs import DatetimeRangeWithTZColumn -from chronify.sqlalchemy.functions import read_database from chronify.time_range_generator_factory import make_time_range_generator from chronify.datetime_range_generator import DatetimeRangeGeneratorExternalTimeZone from chronify.time import LeapDayAdjustmentType @@ -15,31 +16,31 @@ def check_timestamps( - conn: Connection, - table: Table, + backend: IbisBackend, + table: ir.Table, schema: TableSchema, leap_day_adjustment: Optional[LeapDayAdjustmentType] = None, ) -> None: """Performs checks on time series arrays in a table.""" TimeSeriesChecker( - conn, table, schema, leap_day_adjustment=leap_day_adjustment + backend, table, schema, leap_day_adjustment=leap_day_adjustment ).check_timestamps() class TimeSeriesChecker: """Performs checks on time series arrays in a table. Timestamps in the table will be checked against expected timestamps generated from the - TableSchema's time_config. TZ-awareness of the generated timestamps will match that of thetable. + TableSchema's time_config. TZ-awareness of the generated timestamps will match that of the table. """ def __init__( self, - conn: Connection, - table: Table, + backend: IbisBackend, + table: ir.Table, schema: TableSchema, leap_day_adjustment: Optional[LeapDayAdjustmentType] = None, ) -> None: - self._conn = conn + self._backend = backend self._schema = schema self._table = table self._time_generator = make_time_range_generator( @@ -68,10 +69,13 @@ def _check_expected_timestamps_datetime(self) -> int: """For tz-naive or tz-aware time without external time zone column""" expected = self._time_generator.list_timestamps() time_columns = self._time_generator.list_time_columns() - stmt = select(*(self._table.c[x] for x in time_columns)).distinct() + + # Build query with distinct and not null filters + query = self._table.select([self._table[x] for x in time_columns]).distinct() for col in time_columns: - stmt = stmt.where(self._table.c[col].is_not(None)) - df = read_database(stmt, self._conn, self._schema.time_config) + query = query.filter(self._table[col].notnull()) + + df = read_query(self._backend, query, self._schema.time_config) actual = self._time_generator.list_distinct_timestamps_from_dataframe(df) expected = sorted(set(expected)) # drop duplicates for tz-naive prevailing time check_timestamp_lists(actual, expected) @@ -84,10 +88,13 @@ def _check_expected_timestamps_with_external_time_zone(self) -> int: time_columns = self._time_generator.list_time_columns() assert isinstance(self._schema.time_config, DatetimeRangeWithTZColumn) # for mypy time_columns.append(self._schema.time_config.get_time_zone_column()) - stmt = select(*(self._table.c[x] for x in time_columns)).distinct() + + # Build query with distinct and not null filters + query = self._table.select([self._table[x] for x in time_columns]).distinct() for col in time_columns: - stmt = stmt.where(self._table.c[col].is_not(None)) - df = read_database(stmt, self._conn, self._schema.time_config) + query = query.filter(self._table[col].notnull()) + + df = read_query(self._backend, query, self._schema.time_config) actual_dct = self._time_generator.list_distinct_timestamps_by_time_zone_from_dataframe(df) if sorted(expected_dct.keys()) != sorted(actual_dct.keys()): msg = ( @@ -118,16 +125,18 @@ def _check_null_consistency(self) -> None: any_are_null = " OR ".join((f"{x} IS NULL" for x in time_columns)) query_all = f"SELECT COUNT(*) FROM {self._schema.name} WHERE {all_are_null}" query_any = f"SELECT COUNT(*) FROM {self._schema.name} WHERE {any_are_null}" - res_all = self._conn.execute(text(query_all)).fetchone() - assert res_all is not None - res_any = self._conn.execute(text(query_any)).fetchone() - assert res_any is not None - if res_all[0] != res_any[0]: + + res_all = self._backend.execute_sql_to_df(query_all) + count_all = res_all.iloc[0, 0] + res_any = self._backend.execute_sql_to_df(query_any) + count_any = res_any.iloc[0, 0] + + if count_all != count_any: msg = ( "If any time columns have a NULL value for a row, all time columns in that " "row must be NULL. " - f"Row count where all time values are NULL: {res_all[0]}. " - f"Row count where any time values are NULL: {res_any[0]}. " + f"Row count where all time values are NULL: {count_all}. " + f"Row count where any time values are NULL: {count_any}. " ) raise InvalidTable(msg) @@ -195,12 +204,14 @@ def _check_expected_timestamps_by_time_array(self, count: int) -> None: JOIN t2 ON {on_expr} """ - for result in self._conn.execute(text(query)).fetchall(): - distinct_count_by_ta = result[0] - count_by_ta = result[1] + + results = self._backend.execute_sql_to_df(query) + for _, row in results.iterrows(): + distinct_count_by_ta = row.iloc[0] + count_by_ta = row.iloc[1] if has_tz_naive_prevailing and not count_by_ta == count: - id_vals = result[2:] + id_vals = row.iloc[2:].tolist() values = ", ".join( f"{x}={y}" for x, y in zip(self._schema.time_array_id_columns, id_vals) ) @@ -212,7 +223,7 @@ def _check_expected_timestamps_by_time_array(self, count: int) -> None: raise InvalidTable(msg) if not has_tz_naive_prevailing and not count_by_ta == count == distinct_count_by_ta: - id_vals = result[2:] + id_vals = row.iloc[2:].tolist() values = ", ".join( f"{x}={y}" for x, y in zip(self._schema.time_array_id_columns, id_vals) ) diff --git a/src/chronify/time_series_mapper.py b/src/chronify/time_series_mapper.py index a03a466..2523c5a 100644 --- a/src/chronify/time_series_mapper.py +++ b/src/chronify/time_series_mapper.py @@ -1,7 +1,7 @@ from pathlib import Path from typing import Optional -from sqlalchemy import Engine, MetaData +from chronify.ibis.backend import IbisBackend from chronify.models import TableSchema from chronify.time_series_mapper_representative import MapperRepresentativeTimeToDatetime @@ -20,8 +20,7 @@ def map_time( - engine: Engine, - metadata: MetaData, + backend: IbisBackend, from_schema: TableSchema, to_schema: TableSchema, data_adjustment: Optional[TimeBasedDataAdjustment] = None, @@ -35,7 +34,7 @@ def map_time( to_schema.time_config, DatetimeRange ): MapperRepresentativeTimeToDatetime( - engine, metadata, from_schema, to_schema, data_adjustment, wrap_time_allowed + backend, from_schema, to_schema, data_adjustment, wrap_time_allowed ).map_time( scratch_dir=scratch_dir, output_file=output_file, @@ -45,7 +44,7 @@ def map_time( to_schema.time_config, DatetimeRange ): MapperDatetimeToDatetime( - engine, metadata, from_schema, to_schema, data_adjustment, wrap_time_allowed + backend, from_schema, to_schema, data_adjustment, wrap_time_allowed ).map_time( scratch_dir=scratch_dir, output_file=output_file, @@ -55,7 +54,7 @@ def map_time( to_schema.time_config, DatetimeRange ): MapperIndexTimeToDatetime( - engine, metadata, from_schema, to_schema, data_adjustment, wrap_time_allowed + backend, from_schema, to_schema, data_adjustment, wrap_time_allowed ).map_time( scratch_dir=scratch_dir, output_file=output_file, @@ -67,7 +66,7 @@ def map_time( # No way to generate expected timestamps for YearMonthDayPeriodTimeNTZ # Is there a way to only check the output datetime timestamps? MapperColumnRepresentativeToDatetime( - engine, metadata, from_schema, to_schema, data_adjustment, wrap_time_allowed + backend, from_schema, to_schema, data_adjustment, wrap_time_allowed ).map_time( scratch_dir=scratch_dir, output_file=output_file, diff --git a/src/chronify/time_series_mapper_base.py b/src/chronify/time_series_mapper_base.py index 127d9db..981ea76 100644 --- a/src/chronify/time_series_mapper_base.py +++ b/src/chronify/time_series_mapper_base.py @@ -1,24 +1,27 @@ import abc -from functools import reduce -from operator import and_ from pathlib import Path from typing import Any, Optional +import ibis.expr.datatypes as dt import pandas as pd from loguru import logger -from sqlalchemy import Engine, MetaData, Table, select, text, func -from chronify.hive_functions import create_materialized_view -from chronify.sqlalchemy.functions import ( +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import ( create_view_from_parquet, - write_database, - write_query_to_parquet, + write_parquet, + write_table, ) +from chronify.spark_functions import create_materialized_view from chronify.models import TableSchema, MappingTableSchema from chronify.exceptions import ConflictingInputsError, InvalidOperation -from chronify.utils.sqlalchemy_table import create_table from chronify.time_series_checker import check_timestamps -from chronify.time import TimeIntervalType, ResamplingOperationType, AggregationType +from chronify.time import ( + TimeIntervalType, + ResamplingOperationType, + AggregationType, + TimeDataType, +) from chronify.time_configs import TimeBasedDataAdjustment from chronify.utils.path_utils import to_path @@ -28,16 +31,14 @@ class TimeSeriesMapperBase(abc.ABC): def __init__( self, - engine: Engine, - metadata: MetaData, + backend: IbisBackend, from_schema: TableSchema, to_schema: TableSchema, data_adjustment: Optional[TimeBasedDataAdjustment] = None, wrap_time_allowed: bool = False, resampling_operation: Optional[ResamplingOperationType] = None, ) -> None: - self._engine = engine - self._metadata = metadata + self._backend = backend self._from_schema = from_schema self._to_schema = to_schema # data_adjustment is used in mapping creation and time check of mapped time @@ -86,13 +87,46 @@ def map_time(self) -> None: """Convert time columns with from_schema to to_schema configuration.""" +def _ensure_mapping_types_match_source( + df_mapping: pd.DataFrame, + from_schema: TableSchema, + backend: IbisBackend, +) -> pd.DataFrame: + """Ensure mapping DataFrame 'from_*' column types match source table column types. + + When unpivoting, column names become string values. The mapping DataFrame may have + integer types for columns like 'from_hour'. This function casts those columns + to match the source table's types. + """ + source_table = backend.table(from_schema.name) + source_schema = source_table.schema() + + df = df_mapping.copy() + for col in df.columns: + if col.startswith("from_"): + source_col = col.removeprefix("from_") + if source_col in source_schema: + source_type = source_schema[source_col] + # Only coerce types when Ibis can determine the source type + # Skip if source type is unknown (e.g., SQLite datetime) + try: + if source_type.is_string() and not pd.api.types.is_string_dtype(df[col]): + df[col] = df[col].astype(str) + elif source_type.is_integer() and not pd.api.types.is_integer_dtype(df[col]): + df[col] = df[col].astype(int) + except AttributeError: + # Unknown type - skip coercion + pass + + return df + + def apply_mapping( df_mapping: pd.DataFrame, mapping_schema: MappingTableSchema, from_schema: TableSchema, to_schema: TableSchema, - engine: Engine, - metadata: MetaData, + backend: IbisBackend, data_adjustment: TimeBasedDataAdjustment, resampling_operation: Optional[ResamplingOperationType] = None, scratch_dir: Optional[Path] = None, @@ -102,24 +136,38 @@ def apply_mapping( """ Apply mapping to create result table with process to clean up and roll back if checks fail """ - with engine.begin() as conn: - write_database( - df_mapping, - conn, - mapping_schema.name, - mapping_schema.time_configs, - if_table_exists="fail", - scratch_dir=scratch_dir, - ) - metadata.reflect(engine, views=True) + # Ensure mapping DataFrame column types match source table column types + df_mapping = _ensure_mapping_types_match_source(df_mapping, from_schema, backend) + + # Debug: Print mapping DF around target time + try: + debug_ts = pd.Timestamp("2020-11-01 08:00:00", tz="EST") + from_cols = [c for c in df_mapping.columns if c.startswith("from_")] + if from_cols: + debug_row = df_mapping[df_mapping[from_cols[0]] == debug_ts] + print(f"DEBUG: Mapping DF row for {debug_ts}:\n{debug_row}") + else: + print("DEBUG: No from_ columns in mapping DF") + except Exception as e: + print(f"DEBUG: Error printing mapping DF: {e}") + + # Create the mapping table + write_table( + backend, + df_mapping, + mapping_schema.name, + mapping_schema.time_configs, + if_exists="fail", + scratch_dir=scratch_dir, + ) + created_tmp_view = False try: _apply_mapping( mapping_schema.name, from_schema, to_schema, - engine, - metadata, + backend, resampling_operation=resampling_operation, scratch_dir=scratch_dir, output_file=output_file, @@ -127,108 +175,262 @@ def apply_mapping( if check_mapped_timestamps: if output_file is not None: output_file = to_path(output_file) - with engine.begin() as conn: - create_view_from_parquet(conn, to_schema.name, output_file) - metadata.reflect(engine, views=True) + create_view_from_parquet(backend, to_schema.name, output_file) created_tmp_view = True - mapped_table = Table(to_schema.name, metadata) - with engine.connect() as conn: - try: - check_timestamps( - conn, - mapped_table, - to_schema, - leap_day_adjustment=data_adjustment.leap_day_adjustment, - ) - except Exception: - logger.exception( - "check_timestamps failed on mapped table {}. Drop it", - to_schema.name, - ) - if output_file is None: - conn.execute(text(f"DROP TABLE {to_schema.name}")) - raise + mapped_table = backend.table(to_schema.name) + try: + check_timestamps( + backend, + mapped_table, + to_schema, + leap_day_adjustment=data_adjustment.leap_day_adjustment, + ) + except Exception: + logger.exception( + "check_timestamps failed on mapped table {}. Drop it", + to_schema.name, + ) + if output_file is None: + backend.drop_table(to_schema.name, if_exists=True) + raise finally: - with engine.begin() as conn: - table_type = "view" if engine.name == "hive" else "table" - conn.execute(text(f"DROP {table_type} IF EXISTS {mapping_schema.name}")) + table_type = "view" if backend.name == "spark" else "table" + backend.execute_sql(f"DROP {table_type} IF EXISTS {mapping_schema.name}") + + if created_tmp_view: + backend.drop_view(to_schema.name, if_exists=True) + + +def _build_join_predicate( + left_table: Any, + right_table: Any, + key: str, + left_type: Any, + right_type: Any, +) -> Any: + """Build a join predicate with appropriate type casting.""" + left_col = left_table[key] + right_col = right_table[f"from_{key}"] + + if left_type is None or right_type is None: + return left_col == right_col + + left_is_unknown = not hasattr(left_type, "is_timestamp") or str(left_type).startswith( + "unknown" + ) + right_is_timestamp = hasattr(right_type, "is_timestamp") and right_type.is_timestamp() + left_is_timestamp = hasattr(left_type, "is_timestamp") and left_type.is_timestamp() + right_is_string = hasattr(right_type, "is_string") and right_type.is_string() + left_is_string = hasattr(left_type, "is_string") and left_type.is_string() - if created_tmp_view: - conn.execute(text(f"DROP VIEW IF EXISTS {to_schema.name}")) - metadata.remove(Table(to_schema.name, metadata)) + if left_is_unknown and right_is_timestamp: + left_col = left_col.cast("timestamp") + elif left_is_timestamp and right_is_string: + right_col = right_col.cast("timestamp") + elif left_is_string and right_is_timestamp: + left_col = left_col.cast("timestamp") - metadata.remove(Table(mapping_schema.name, metadata)) - metadata.reflect(engine, views=True) + return left_col == right_col + + +def _get_timestamp_target_dtype(to_schema: TableSchema, backend_name: str) -> Any: + """Get the target dtype for timestamp column casting.""" + time_config = to_schema.time_config + dtype = getattr(time_config, "dtype", None) + if dtype == TimeDataType.TIMESTAMP_TZ: + start = getattr(time_config, "start", None) + tz_obj = start.tzinfo if start is not None else None + if backend_name == "sqlite": + tz_str = "UTC" + elif tz_obj is not None and hasattr(tz_obj, "key"): + tz_str = tz_obj.key + elif tz_obj: + tz_str = str(tz_obj) + else: + tz_str = "UTC" + return dt.Timestamp(timezone=tz_str) + return dt.timestamp + + +def _build_column_expr( + joined: Any, + col: str, + left_table_columns: list[str], + joined_columns: list[str], +) -> Any: + """Build a column expression, handling name conflicts from join.""" + if col in left_table_columns and f"{col}_right" in joined_columns: + return joined[f"{col}_right"].name(col) + return joined[col] def _apply_mapping( mapping_table_name: str, from_schema: TableSchema, to_schema: TableSchema, - engine: Engine, - metadata: MetaData, + backend: IbisBackend, resampling_operation: Optional[ResamplingOperationType] = None, scratch_dir: Optional[Path] = None, output_file: Optional[Path] = None, ) -> None: - """Apply mapping to create result as a table according to_schema - - Columns used to join the from_table are prefixed with "from_" in the mapping table + """Apply mapping to create result as a table according to_schema. + + Columns used to join the from_table are prefixed with "from_" in the mapping table. """ - left_table = Table(from_schema.name, metadata) - right_table = Table(mapping_table_name, metadata) - left_table_columns = [x.name for x in left_table.columns] - right_table_columns = [x.name for x in right_table.columns] + left_table = backend.table(from_schema.name) + right_table = backend.table(mapping_table_name) + + left_table_columns = list(left_table.columns) + right_table_columns = list(right_table.columns) left_table_pass_thru_columns = set(left_table_columns).difference( set(from_schema.list_columns()) ) - val_col = to_schema.value_column # from left_table + val_col = to_schema.value_column final_cols = set(to_schema.list_columns()).union(left_table_pass_thru_columns) right_cols = set(right_table_columns).intersection(final_cols) left_cols = final_cols - right_cols - {val_col} - select_stmt: list[Any] = [left_table.c[x] for x in left_cols] - select_stmt += [right_table.c[x] for x in right_cols] - - tval_col = left_table.c[val_col] - if "factor" in right_table_columns: - tval_col *= right_table.c["factor"] # type: ignore - if not resampling_operation: - select_stmt.append(tval_col) - else: - groupby_stmt = select_stmt.copy() - match resampling_operation: - case AggregationType.SUM: - select_stmt.append(func.sum(tval_col).label(val_col)) - # case AggregationType.AVG: - # select_stmt.append(func.avg(tval_col).label(val_col)) - # case AggregationType.MIN: - # select_stmt.append(func.min(tval_col).label(val_col)) - # case AggregationType.MAX: - # select_stmt.append(func.max(tval_col).label(val_col)) - case _: - msg = f"Unsupported {resampling_operation=}" - raise InvalidOperation(msg) - + # Build join predicates from_keys = [x for x in right_table_columns if x.startswith("from_")] keys = [x.removeprefix("from_") for x in from_keys] assert set(keys).issubset( set(left_table_columns) ), f"Keys {keys} not in table={from_schema.name}" - on_stmt = reduce(and_, (left_table.c[x] == right_table.c["from_" + x] for x in keys)) - query = select(*select_stmt).select_from(left_table).join(right_table, on_stmt) - if resampling_operation: - query = query.group_by(*groupby_stmt) + predicates = _build_join_predicates(left_table, right_table, keys) + joined = left_table.join(right_table, predicates) + joined_columns = list(joined.columns) + + # Build select columns + select_cols = _build_select_columns( + joined, left_cols, right_cols, left_table_columns, joined_columns, to_schema, backend.name + ) + + # Handle value column with optional factor multiplication + tval_col = joined[val_col] + if "factor" in right_table_columns: + tval_col = tval_col * joined["factor"] + + query = _build_query( + joined, + select_cols, + tval_col, + val_col, + left_cols, + right_cols, + left_table_columns, + joined_columns, + resampling_operation, + ) + + _write_result(query, to_schema.name, backend, scratch_dir, output_file) + + +def _build_join_predicates(left_table: Any, right_table: Any, keys: list[str]) -> list[Any]: + """Build join predicates with type mismatch handling.""" + left_schema = left_table.schema() + right_schema = right_table.schema() + predicates = [] + + for k in keys: + left_type = left_schema.get(k) + right_type = right_schema.get(f"from_{k}") + try: + pred = _build_join_predicate(left_table, right_table, k, left_type, right_type) + predicates.append(pred) + except Exception: + # Fallback: cast both to string + predicates.append( + left_table[k].cast("string") == right_table[f"from_{k}"].cast("string") + ) + + return predicates + + +def _build_select_columns( + joined: Any, + left_cols: set[str], + right_cols: set[str], + left_table_columns: list[str], + joined_columns: list[str], + to_schema: TableSchema, + backend_name: str, +) -> list[Any]: + """Build the list of columns to select from the joined table.""" + select_cols: list[Any] = [] + + # Left table columns + for col in left_cols: + select_cols.append(joined[col]) + + # Right table columns with potential name conflict handling + time_column = getattr(to_schema.time_config, "time_column", None) + dtype = getattr(to_schema.time_config, "dtype", None) + needs_timestamp_cast = dtype is not None and dtype in ( + TimeDataType.TIMESTAMP_TZ, + TimeDataType.TIMESTAMP_NTZ, + ) + + for col in right_cols: + col_expr = _build_column_expr(joined, col, left_table_columns, joined_columns) + + if col == time_column and needs_timestamp_cast: + target_dtype = _get_timestamp_target_dtype(to_schema, backend_name) + col_expr = col_expr.cast(target_dtype).name(col) + + select_cols.append(col_expr) + + return select_cols + +def _build_query( + joined: Any, + select_cols: list[Any], + tval_col: Any, + val_col: str, + left_cols: set[str], + right_cols: set[str], + left_table_columns: list[str], + joined_columns: list[str], + resampling_operation: Optional[ResamplingOperationType], +) -> Any: + """Build the final query with optional aggregation.""" + if not resampling_operation: + select_cols.append(tval_col.name(val_col)) + return joined.select(select_cols) + + # Aggregation case + groupby_cols = [joined[col] for col in left_cols] + for col in right_cols: + groupby_cols.append(_build_column_expr(joined, col, left_table_columns, joined_columns)) + + match resampling_operation: + case AggregationType.SUM: + agg_col = tval_col.sum().name(val_col) + case _: + msg = f"Unsupported {resampling_operation=}" + raise InvalidOperation(msg) + + return joined.group_by(groupby_cols).aggregate(agg_col) + + +def _write_result( + query: Any, + table_name: str, + backend: IbisBackend, + scratch_dir: Optional[Path], + output_file: Optional[Path], +) -> None: + """Write the query result to the appropriate destination.""" if output_file is not None: output_file = to_path(output_file) - write_query_to_parquet(engine, str(query), output_file, overwrite=True) + write_parquet(backend, query, output_file, overwrite=True) return - if engine.name == "hive": + if backend.name == "spark": create_materialized_view( - str(query), to_schema.name, engine, metadata, scratch_dir=scratch_dir + str(query.compile()), table_name, backend, scratch_dir=scratch_dir ) else: - create_table(to_schema.name, query, engine, metadata) + backend.create_table(table_name, query) diff --git a/src/chronify/time_series_mapper_column_representative_to_datetime.py b/src/chronify/time_series_mapper_column_representative_to_datetime.py index ed549de..506d2b2 100644 --- a/src/chronify/time_series_mapper_column_representative_to_datetime.py +++ b/src/chronify/time_series_mapper_column_representative_to_datetime.py @@ -1,10 +1,11 @@ from typing import Optional, Generator import re -import sqlalchemy as sa from pathlib import Path import pandas as pd from datetime import datetime +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import read_query, write_table from chronify.exceptions import InvalidParameter, InvalidValue from chronify.time_series_mapper_base import TimeSeriesMapperBase, apply_mapping from chronify.time_configs import ( @@ -18,8 +19,6 @@ ) from chronify.datetime_range_generator import DatetimeRangeGenerator from chronify.models import MappingTableSchema, TableSchema -from chronify.sqlalchemy.functions import read_database, write_database -from chronify.utils.sqlalchemy_table import create_table class MapperColumnRepresentativeToDatetime(TimeSeriesMapperBase): @@ -51,16 +50,13 @@ class MapperColumnRepresentativeToDatetime(TimeSeriesMapperBase): def __init__( self, - engine: sa.Engine, - metadata: sa.MetaData, + backend: IbisBackend, from_schema: TableSchema, to_schema: TableSchema, data_adjustment: Optional[TimeBasedDataAdjustment] = None, wrap_time_allowed: bool = False, ) -> None: - super().__init__( - engine, metadata, from_schema, to_schema, data_adjustment, wrap_time_allowed - ) + super().__init__(backend, from_schema, to_schema, data_adjustment, wrap_time_allowed) if not isinstance(to_schema.time_config, DatetimeRange): msg = "Target schema does not have DatetimeRange time config. Use a different mapper." @@ -102,8 +98,7 @@ def map_time( mapping_schema, from_schema, self._to_schema, - self._engine, - self._metadata, + self._backend, self._data_adjustment, scratch_dir=scratch_dir, output_file=output_file, @@ -111,9 +106,8 @@ def map_time( ) if drop_table: - with self._engine.begin() as conn: - table_type = "view" if self._engine.name == "hive" else "table" - conn.execute(sa.text(f"DROP {table_type} IF EXISTS {drop_table}")) + table_type = "view" if self._backend.name == "spark" else "table" + self._backend.execute_sql(f"DROP {table_type} IF EXISTS {drop_table}") def check_schema_consistency(self) -> None: if isinstance(self._from_time_config, MonthDayHourTimeNTZ): @@ -137,36 +131,39 @@ def _intermediate_mapping_ymdp_to_ymdh(self, scratch_dir: Path | None) -> TableS """Convert ymdp to ymdh for intermediate mapping.""" mapping_table_name = "intermediate_ymdp_to_ymdh" period_col = self._from_time_config.hour_columns[0] - with self._engine.begin() as conn: - periods = read_database( - f"SELECT DISTINCT {period_col} FROM {self._from_schema.name}", - conn, - self._from_time_config, - ) - df_mapping = generate_period_mapping(periods.iloc[:, 0]) - write_database( - df_mapping, - conn, - mapping_table_name, - [self._from_time_config], - if_table_exists="replace", - scratch_dir=scratch_dir, - ) - self._metadata.reflect(self._engine) - ymdp_table = sa.Table(self._from_schema.name, self._metadata) - mapping_table = sa.Table(mapping_table_name, self._metadata) + # Read distinct periods + table = self._backend.table(self._from_schema.name) + query = table.select(table[period_col]).distinct() + periods = read_query(self._backend, query, self._from_time_config) + df_mapping = generate_period_mapping(periods.iloc[:, 0]) + + write_table( + self._backend, + df_mapping, + mapping_table_name, + [self._from_time_config], + if_exists="replace", + scratch_dir=scratch_dir, + ) + + # Build join query using Ibis + ymdp_table = self._backend.table(self._from_schema.name) + mapping_table = self._backend.table(mapping_table_name) + + # Select all columns from ymdp_table except period_col, plus hour from mapping + ymdp_cols = list(ymdp_table.columns) + select_cols = [ymdp_table[col] for col in ymdp_cols if col != period_col] + select_cols.append(mapping_table["hour"]) - select_statement = [col for col in ymdp_table.columns if col.name != period_col] - select_statement.append(mapping_table.c["hour"]) - query = ( - sa.select(*select_statement) - .select_from(ymdp_table) - .join(mapping_table, ymdp_table.c[period_col] == mapping_table.c["from_period"]) + # Join and select + joined = ymdp_table.join( + mapping_table, ymdp_table[period_col] == mapping_table["from_period"] ) + query = joined.select(select_cols) intermediate_ymdh_table_name = "intermediate_Ymdh" - create_table(intermediate_ymdh_table_name, query, self._engine, self._metadata) + self._backend.create_table(intermediate_ymdh_table_name, query) assert isinstance( self._from_time_config, YearMonthDayPeriodTimeNTZ diff --git a/src/chronify/time_series_mapper_datetime.py b/src/chronify/time_series_mapper_datetime.py index 108aaff..b75f632 100644 --- a/src/chronify/time_series_mapper_datetime.py +++ b/src/chronify/time_series_mapper_datetime.py @@ -3,8 +3,8 @@ from typing import Optional import pandas as pd -from sqlalchemy import Engine, MetaData +from chronify.ibis.backend import IbisBackend from chronify.models import TableSchema, MappingTableSchema from chronify.exceptions import InvalidParameter, ConflictingInputsError from chronify.time_series_mapper_base import TimeSeriesMapperBase, apply_mapping @@ -22,16 +22,13 @@ class MapperDatetimeToDatetime(TimeSeriesMapperBase): def __init__( self, - engine: Engine, - metadata: MetaData, + backend: IbisBackend, from_schema: TableSchema, to_schema: TableSchema, data_adjustment: Optional[TimeBasedDataAdjustment] = None, wrap_time_allowed: bool = False, ) -> None: - super().__init__( - engine, metadata, from_schema, to_schema, data_adjustment, wrap_time_allowed - ) + super().__init__(backend, from_schema, to_schema, data_adjustment, wrap_time_allowed) if self._from_schema == self._to_schema and self._data_adjustment is None: msg = f"from_schema is the same as to_schema and no data_adjustment, nothing to do.\n{self._from_schema}" logger.info(msg) @@ -77,8 +74,7 @@ def map_time( mapping_schema, self._from_schema, self._to_schema, - self._engine, - self._metadata, + self._backend, self._data_adjustment, scratch_dir=scratch_dir, output_file=output_file, @@ -106,17 +102,12 @@ def _create_mapping(self) -> tuple[pd.DataFrame, MappingTableSchema]: # get standard time zone of to_tz to_tz_std = get_standard_time_zone(to_tz) # tz-naive time does not have skips/dups, so always localize in std tz first - # mypy complains with error: "Properties" has no attribute "tz_localize" - # but the type is pandas.core.indexes.accessors.DatetimeProperties and the - # attribute exists. - ser_from = ser_from.dt.tz_localize(to_tz_std).dt.tz_convert(to_tz) # type: ignore - pass + ser_from = ser_from.dt.tz_localize(to_tz_std).dt.tz_convert(to_tz) case (False, True): # get standard time zone of fm_tz fm_tz_std = get_standard_time_zone(fm_tz) # convert to standard time zone of fm_tz before making it tz-naive - ser_from = ser_from.dt.tz_convert(fm_tz_std).dt.tz_localize(to_tz) # type: ignore - pass + ser_from = ser_from.dt.tz_convert(fm_tz_std).dt.tz_localize(to_tz) match (self._adjust_interval, self._wrap_time_allowed): case (True, _): ser = pd.Series( diff --git a/src/chronify/time_series_mapper_index_time.py b/src/chronify/time_series_mapper_index_time.py index 4ef3925..79b014c 100644 --- a/src/chronify/time_series_mapper_index_time.py +++ b/src/chronify/time_series_mapper_index_time.py @@ -5,8 +5,9 @@ from datetime import datetime, timedelta import pandas as pd -from sqlalchemy import Engine, MetaData, Table, select +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import read_query from chronify.models import TableSchema, MappingTableSchema from chronify.exceptions import InvalidParameter, ConflictingInputsError from chronify.time_series_mapper_base import TimeSeriesMapperBase, apply_mapping @@ -20,7 +21,6 @@ from chronify.time_range_generator_factory import make_time_range_generator from chronify.time_series_mapper_datetime import MapperDatetimeToDatetime from chronify.time import TimeDataType, TimeType, DaylightSavingAdjustmentType, AggregationType -from chronify.sqlalchemy.functions import read_database logger = logging.getLogger(__name__) @@ -28,17 +28,14 @@ class MapperIndexTimeToDatetime(TimeSeriesMapperBase): def __init__( self, - engine: Engine, - metadata: MetaData, + backend: IbisBackend, from_schema: TableSchema, to_schema: TableSchema, data_adjustment: Optional[TimeBasedDataAdjustment] = None, wrap_time_allowed: bool = False, ) -> None: # TODO: refactor to use new time configs - super().__init__( - engine, metadata, from_schema, to_schema, data_adjustment, wrap_time_allowed - ) + super().__init__(backend, from_schema, to_schema, data_adjustment, wrap_time_allowed) self._dst_adjustment = self._data_adjustment.daylight_saving_adjustment if not isinstance(self._from_schema.time_config, IndexTimeRangeBase): msg = "Source schema does not have IndexTimeRangeBase time config. Use a different mapper." @@ -119,8 +116,7 @@ def map_time( mapping_schema, self._from_schema, mapped_schema, - self._engine, - self._metadata, + self._backend, TimeBasedDataAdjustment(), resampling_operation=resampling_operation, scratch_dir=scratch_dir, @@ -129,8 +125,7 @@ def map_time( # Convert from represented datetime to dst time_config MapperDatetimeToDatetime( - self._engine, - self._metadata, + self._backend, mapped_schema, self._to_schema, self._data_adjustment, @@ -219,10 +214,9 @@ def _create_interm_map_with_time_zone( assert tz_col is not None, "Expecting a time zone column for INDEX_TZ_COL" from_tz_col = "from_" + tz_col - with self._engine.connect() as conn: - table = Table(self._from_schema.name, self._metadata) - stmt = select(table.c[tz_col]).distinct().where(table.c[tz_col].is_not(None)) - time_zones = read_database(stmt, conn, self._from_time_config)[tz_col].to_list() + table = self._backend.table(self._from_schema.name) + query = table.select(table[tz_col]).distinct().filter(table[tz_col].notnull()) + time_zones = read_query(self._backend, query, self._from_time_config)[tz_col].to_list() from_time_config = self._from_time_config.model_copy( update={"time_column": from_time_col, "time_zone_column": from_tz_col} @@ -290,10 +284,9 @@ def _create_interm_map_with_time_zone_and_dst_adjustment( assert tz_col is not None, "Expecting a time zone column for INDEX_TZ_COL" from_tz_col = "from_" + tz_col - with self._engine.connect() as conn: - table = Table(self._from_schema.name, self._metadata) - stmt = select(table.c[tz_col]).distinct().where(table.c[tz_col].is_not(None)) - time_zones = read_database(stmt, conn, self._from_time_config)[tz_col].to_list() + table = self._backend.table(self._from_schema.name) + query = table.select(table[tz_col]).distinct().filter(table[tz_col].notnull()) + time_zones = read_query(self._backend, query, self._from_time_config)[tz_col].to_list() from_time_config = self._from_time_config.model_copy( update={"time_column": from_time_col, "time_zone_column": from_tz_col} diff --git a/src/chronify/time_series_mapper_representative.py b/src/chronify/time_series_mapper_representative.py index 9c545cd..409b7c8 100644 --- a/src/chronify/time_series_mapper_representative.py +++ b/src/chronify/time_series_mapper_representative.py @@ -3,9 +3,9 @@ from typing import Optional import pandas as pd -from sqlalchemy import Engine, MetaData, Table, select -from chronify.sqlalchemy.functions import read_database +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import read_query from chronify.models import TableSchema, MappingTableSchema from chronify.exceptions import ( InvalidParameter, @@ -27,16 +27,13 @@ class MapperRepresentativeTimeToDatetime(TimeSeriesMapperBase): def __init__( self, - engine: Engine, - metadata: MetaData, + backend: IbisBackend, from_schema: TableSchema, to_schema: TableSchema, data_adjustment: Optional[TimeBasedDataAdjustment] = None, wrap_time_allowed: bool = False, ) -> None: - super().__init__( - engine, metadata, from_schema, to_schema, data_adjustment, wrap_time_allowed - ) + super().__init__(backend, from_schema, to_schema, data_adjustment, wrap_time_allowed) if not isinstance(from_schema.time_config, RepresentativePeriodTimeBase): msg = "source schema does not have RepresentativePeriodTimeBase time config. Use a different mapper." raise InvalidParameter(msg) @@ -69,8 +66,7 @@ def map_time( mapping_schema, self._from_schema, self._to_schema, - self._engine, - self._metadata, + self._backend, self._data_adjustment, scratch_dir=scratch_dir, output_file=output_file, @@ -107,10 +103,9 @@ def _create_mapping(self, is_tz_naive: bool) -> tuple[pd.DataFrame, MappingTable else: tz_col = self._from_time_config.get_time_zone_column() assert tz_col is not None, "Expecting a time zone column for REPRESENTATIVE time" - with self._engine.connect() as conn: - table = Table(self._from_schema.name, self._metadata) - stmt = select(table.c[tz_col]).distinct().where(table.c[tz_col].is_not(None)) - time_zones = read_database(stmt, conn, self._from_time_config)[tz_col].to_list() + table = self._backend.table(self._from_schema.name) + query = table.select(table[tz_col]).distinct().filter(table[tz_col].notnull()) + time_zones = read_query(self._backend, query, self._from_time_config)[tz_col].to_list() df = self._generator.create_tz_aware_mapping_dataframe( dft, time_col, time_zones, tz_col ) diff --git a/src/chronify/time_utils.py b/src/chronify/time_utils.py index 0744c46..ea5cb4d 100644 --- a/src/chronify/time_utils.py +++ b/src/chronify/time_utils.py @@ -1,16 +1,16 @@ """Functions related to time""" import logging -from numpy.typing import NDArray -import numpy as np +from collections.abc import Sequence from datetime import datetime, timedelta, timezone, tzinfo -from zoneinfo import ZoneInfo, ZoneInfoNotFoundError + +import numpy as np import pandas as pd +from numpy.typing import NDArray +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError -from chronify.time import ( - TimeIntervalType, -) from chronify.exceptions import InvalidParameter +from chronify.time import TimeIntervalType logger = logging.getLogger(__name__) @@ -30,7 +30,7 @@ def adjust_timestamp_by_dst_offset(timestamp: datetime, resolution: timedelta) - def shifted_interval_timestamps( - ts_list: list[datetime], + ts_list: Sequence[datetime], from_interval_type: TimeIntervalType, to_interval_type: TimeIntervalType, ) -> list[datetime]: @@ -73,8 +73,8 @@ def shifted_interval_timestamps( def wrapped_time_timestamps( - ts_list: list[datetime], - to_timestamps: list[datetime], + ts_list: Sequence[datetime], + to_timestamps: Sequence[datetime], ) -> list[datetime]: """Returns the replacement timestamps in order to wrap the ts_list into the to_timestamps range. @@ -108,10 +108,10 @@ def wrapped_time_timestamps( def rolled_interval_timestamps( - ts_list: list[datetime], + ts_list: Sequence[datetime], from_interval_type: TimeIntervalType, to_interval_type: TimeIntervalType, - to_timestamps: list[datetime], + to_timestamps: Sequence[datetime], ) -> list[datetime]: """Roll ts_list by shifting time interval based on interval type and then wrapping timestamps according to to_timestamps. diff --git a/src/chronify/time_zone_converter.py b/src/chronify/time_zone_converter.py index a686e80..c0f1503 100644 --- a/src/chronify/time_zone_converter.py +++ b/src/chronify/time_zone_converter.py @@ -1,11 +1,12 @@ import abc from zoneinfo import ZoneInfo from datetime import datetime, tzinfo -from sqlalchemy import Engine, MetaData, Table, select from typing import Optional from pathlib import Path import pandas as pd +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import read_query from chronify.models import TableSchema, MappingTableSchema from chronify.time_configs import ( DatetimeRangeBase, @@ -19,15 +20,13 @@ from chronify.exceptions import InvalidParameter, MissingValue from chronify.time_series_mapper_base import apply_mapping from chronify.time_range_generator_factory import make_time_range_generator -from chronify.sqlalchemy.functions import read_database from chronify.time import TimeDataType, TimeType from chronify.time_utils import wrapped_time_timestamps, get_tzname # TODO - retain original timestamp column def convert_time_zone( - engine: Engine, - metadata: MetaData, + backend: IbisBackend, src_schema: TableSchema, to_time_zone: tzinfo | None, scratch_dir: Optional[Path] = None, @@ -40,10 +39,8 @@ def convert_time_zone( Parameters ---------- - engine : sqlalchemy.Engine - SQLAlchemy engine. - metadata : sqlalchemy.MetaData - SQLAlchemy metadata. + backend : IbisBackend + Ibis backend. src_schema : TableSchema Defines the source table in the database. to_time_zone : tzinfo or None @@ -61,7 +58,7 @@ def convert_time_zone( TableSchema Schema of output table with converted timestamps. """ - tzc = TimeZoneConverter(engine, metadata, src_schema, to_time_zone) + tzc = TimeZoneConverter(backend, src_schema, to_time_zone) tzc.convert_time_zone( scratch_dir=scratch_dir, output_file=output_file, @@ -72,8 +69,7 @@ def convert_time_zone( def convert_time_zone_by_column( - engine: Engine, - metadata: MetaData, + backend: IbisBackend, src_schema: TableSchema, time_zone_column: str, wrap_time_allowed: Optional[bool] = False, @@ -86,10 +82,8 @@ def convert_time_zone_by_column( Parameters ---------- - engine : sqlalchemy.Engine - sqlalchemy engine - metadata : sqlalchemy.MetaData - sqlalchemy metadata + backend : IbisBackend + Ibis backend src_schema : TableSchema Defines the source table in the database. time_zone_column : str @@ -116,9 +110,7 @@ def convert_time_zone_by_column( dst_schema : TableSchema schema of output table with converted timestamps """ - tzc = TimeZoneConverterByColumn( - engine, metadata, src_schema, time_zone_column, wrap_time_allowed - ) + tzc = TimeZoneConverterByColumn(backend, src_schema, time_zone_column, wrap_time_allowed) tzc.convert_time_zone( scratch_dir=scratch_dir, output_file=output_file, @@ -132,19 +124,18 @@ class TimeZoneConverterBase(abc.ABC): def __init__( self, - engine: Engine, - metadata: MetaData, + backend: IbisBackend, from_schema: TableSchema, ): - self._engine = engine - self._metadata = metadata + self._backend = backend self._check_from_schema(from_schema) self._from_schema = from_schema def _check_from_schema(self, from_schema: TableSchema) -> None: - msg = "" if not isinstance(from_schema.time_config, DatetimeRange): - msg += "Source schema must have DatetimeRange time config. " + msg = "Source schema must have DatetimeRange time config. " + raise InvalidParameter(msg) + msg = "" if from_schema.time_config.dtype != TimeDataType.TIMESTAMP_TZ: msg += "Source schema time config dtype must be Timestamp_TZ. " if from_schema.time_config.start_time_is_tz_naive(): @@ -187,12 +178,11 @@ class TimeZoneConverter(TimeZoneConverterBase): def __init__( self, - engine: Engine, - metadata: MetaData, + backend: IbisBackend, from_schema: TableSchema, to_time_zone: tzinfo | None, ): - super().__init__(engine, metadata, from_schema) + super().__init__(backend, from_schema) self._to_time_zone = to_time_zone self._to_schema = self.generate_to_schema() @@ -219,7 +209,7 @@ def generate_to_time_config(self) -> DatetimeRangeWithTZColumn: def generate_to_schema(self) -> TableSchema: to_time_config = self.generate_to_time_config() - id_cols = self._from_schema.time_array_id_columns + id_cols = self._from_schema.time_array_id_columns.copy() if to_time_config.time_zone_column not in id_cols: id_cols.append(to_time_config.time_zone_column) to_schema: TableSchema = self._from_schema.model_copy( @@ -244,8 +234,7 @@ def convert_time_zone( mapping_schema, self._from_schema, self._to_schema, - self._engine, - self._metadata, + self._backend, TimeBasedDataAdjustment(), scratch_dir=scratch_dir, output_file=output_file, @@ -319,8 +308,7 @@ class TimeZoneConverterByColumn(TimeZoneConverterBase): def __init__( self, - engine: Engine, - metadata: MetaData, + backend: IbisBackend, from_schema: TableSchema, time_zone_column: str, wrap_time_allowed: Optional[bool] = False, @@ -328,7 +316,7 @@ def __init__( if time_zone_column not in from_schema.time_array_id_columns: msg = f"{time_zone_column=} is missing from {from_schema.time_array_id_columns=}" raise MissingValue(msg) - super().__init__(engine, metadata, from_schema) + super().__init__(backend, from_schema) self.time_zone_column = time_zone_column self._wrap_time_allowed = wrap_time_allowed self._to_schema = self.generate_to_schema() @@ -352,9 +340,9 @@ def generate_to_time_config(self) -> DatetimeRangeBase: return DatetimeRangeWithTZColumn(**time_kwargs) def generate_to_schema(self) -> TableSchema: - id_cols = self._from_schema.time_array_id_columns - if "time_zone" not in id_cols: - id_cols.append("time_zone") + id_cols = self._from_schema.time_array_id_columns.copy() + if self.time_zone_column not in id_cols: + id_cols.append(self.time_zone_column) to_schema: TableSchema = self._from_schema.model_copy( update={ "name": f"{self._from_schema.name}_tz_converted", @@ -377,8 +365,7 @@ def convert_time_zone( mapping_schema, self._from_schema, self._to_schema, - self._engine, - self._metadata, + self._backend, TimeBasedDataAdjustment(), scratch_dir=scratch_dir, output_file=output_file, @@ -386,16 +373,15 @@ def convert_time_zone( ) def _get_time_zones(self) -> list[tzinfo | None]: - with self._engine.connect() as conn: - table = Table(self._from_schema.name, self._metadata) - stmt = ( - select(table.c[self.time_zone_column]) - .distinct() - .where(table.c[self.time_zone_column].is_not(None)) - ) - time_zones = read_database(stmt, conn, self._from_schema.time_config)[ - self.time_zone_column - ].to_list() + table = self._backend.table(self._from_schema.name) + query = ( + table.select(table[self.time_zone_column]) + .distinct() + .filter(table[self.time_zone_column].notnull()) + ) + time_zones = read_query(self._backend, query, self._from_schema.time_config)[ + self.time_zone_column + ].to_list() time_zones = [None if tz == "None" else ZoneInfo(tz) for tz in time_zones] return time_zones diff --git a/src/chronify/time_zone_localizer.py b/src/chronify/time_zone_localizer.py index 721eb56..cdc25f2 100644 --- a/src/chronify/time_zone_localizer.py +++ b/src/chronify/time_zone_localizer.py @@ -1,12 +1,14 @@ import abc +import warnings from zoneinfo import ZoneInfo from datetime import tzinfo -from sqlalchemy import Engine, MetaData, Table, select from typing import Optional from pathlib import Path import pandas as pd from pandas import DatetimeTZDtype +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import read_query from chronify.models import TableSchema, MappingTableSchema from chronify.time_configs import ( DatetimeRangeBase, @@ -21,15 +23,13 @@ from chronify.exceptions import InvalidParameter, MissingValue from chronify.time_series_mapper_base import apply_mapping from chronify.time_range_generator_factory import make_time_range_generator -from chronify.sqlalchemy.functions import read_database from chronify.time import TimeDataType, TimeType from chronify.time_series_mapper import map_time from chronify.time_utils import get_standard_time_zone, is_standard_time_zone def localize_time_zone( - engine: Engine, - metadata: MetaData, + backend: IbisBackend, src_schema: TableSchema, to_time_zone: tzinfo | None, scratch_dir: Optional[Path] = None, @@ -44,10 +44,8 @@ def localize_time_zone( Parameters ---------- - engine : sqlalchemy.Engine - SQLAlchemy engine. - metadata : sqlalchemy.MetaData - SQLAlchemy metadata. + backend : IbisBackend + Ibis backend. src_schema : TableSchema Defines the source table in the database. to_time_zone : tzinfo or None @@ -65,7 +63,7 @@ def localize_time_zone( TableSchema Schema of output table with converted timestamps. """ - tzl = TimeZoneLocalizer(engine, metadata, src_schema, to_time_zone) + tzl = TimeZoneLocalizer(backend, src_schema, to_time_zone) tzl.localize_time_zone( scratch_dir=scratch_dir, output_file=output_file, @@ -76,8 +74,7 @@ def localize_time_zone( def localize_time_zone_by_column( - engine: Engine, - metadata: MetaData, + backend: IbisBackend, src_schema: TableSchema, time_zone_column: Optional[str] = None, scratch_dir: Optional[Path] = None, @@ -89,10 +86,8 @@ def localize_time_zone_by_column( Parameters ---------- - engine : sqlalchemy.Engine - SQLAlchemy engine. - metadata : sqlalchemy.MetaData - sqlalchemy metadata + backend : IbisBackend + Ibis backend. src_schema : TableSchema Defines the source table in the database. time_zone_column : Optional[str] @@ -119,9 +114,7 @@ def localize_time_zone_by_column( ) raise MissingValue(msg) - tzl = TimeZoneLocalizerByColumn( - engine, metadata, src_schema, time_zone_column=time_zone_column - ) + tzl = TimeZoneLocalizerByColumn(backend, src_schema, time_zone_column=time_zone_column) tzl.localize_time_zone( scratch_dir=scratch_dir, output_file=output_file, @@ -135,12 +128,10 @@ class TimeZoneLocalizerBase(abc.ABC): def __init__( self, - engine: Engine, - metadata: MetaData, + backend: IbisBackend, from_schema: TableSchema, ): - self._engine = engine - self._metadata = metadata + self._backend = backend self._from_schema = from_schema @staticmethod @@ -174,27 +165,27 @@ class TimeZoneLocalizer(TimeZoneLocalizerBase): def __init__( self, - engine: Engine, - metadata: MetaData, + backend: IbisBackend, from_schema: TableSchema, to_time_zone: tzinfo | None, ): self._check_from_schema(from_schema) - super().__init__(engine, metadata, from_schema) + super().__init__(backend, from_schema) self._to_time_zone = self._check_standard_time_zone(to_time_zone) self._to_schema = self.generate_to_schema() @staticmethod def _check_from_schema(from_schema: TableSchema) -> None: - msg = "" - if not isinstance(from_schema.time_config, DatetimeRange): - msg += "Source schema does not have DatetimeRange time config. " if isinstance(from_schema.time_config, DatetimeRangeWithTZColumn): - msg += ( - "Instead, time config is of type DatetimeRangeWithTZColumn, " + msg = ( + "Source schema time config is of type DatetimeRangeWithTZColumn, " f"try using TimeZoneLocalizerByColumn(). {from_schema.time_config}" ) raise InvalidParameter(msg) + if not isinstance(from_schema.time_config, DatetimeRange): + msg = f"Source schema does not have DatetimeRange time config.\n{from_schema.time_config}" + raise InvalidParameter(msg) + msg = "" if from_schema.time_config.dtype != TimeDataType.TIMESTAMP_NTZ: msg += "Source schema time config dtype must be TIMESTAMP_NTZ. " if from_schema.time_config.start_time_is_tz_naive() is False: @@ -250,8 +241,7 @@ def localize_time_zone( check_mapped_timestamps: bool = False, ) -> None: map_time( - engine=self._engine, - metadata=self._metadata, + backend=self._backend, from_schema=self._from_schema, to_schema=self._to_schema, scratch_dir=scratch_dir, @@ -288,18 +278,20 @@ class TimeZoneLocalizerByColumn(TimeZoneLocalizerBase): def __init__( self, - engine: Engine, - metadata: MetaData, + backend: IbisBackend, from_schema: TableSchema, time_zone_column: Optional[str] = None, ): self._check_from_schema(from_schema) self._check_time_zone_column(from_schema, time_zone_column) - super().__init__(engine, metadata, from_schema) + # Make a deep copy to avoid mutating the original schema + from_schema_copy = from_schema.model_copy(deep=True) + super().__init__(backend, from_schema_copy) if isinstance(self._from_schema.time_config, DatetimeRange): self.time_zone_column = time_zone_column self._convert_from_time_config_to_datetime_range_with_tz_column() else: + assert isinstance(self._from_schema.time_config, DatetimeRangeWithTZColumn) # mypy self.time_zone_column = self._from_schema.time_config.time_zone_column self._check_standard_time_zones() self._to_schema = self.generate_to_schema() @@ -311,6 +303,8 @@ def _check_from_schema(from_schema: TableSchema) -> None: msg += ( "Source schema must have DatetimeRange or DatetimeRangeWithTZColumn time config. " ) + msg += f"\n{from_schema.time_config}" + raise InvalidParameter(msg) if from_schema.time_config.dtype != TimeDataType.TIMESTAMP_NTZ: msg += "Source schema time config dtype must be TIMESTAMP_NTZ. " if msg != "": @@ -324,7 +318,7 @@ def _check_time_zone_column(from_schema: TableSchema, time_zone_column: Optional and time_zone_column is not None ): msg = f"Input {time_zone_column=} will be ignored. time_zone_column is already defined in the time_config." - raise Warning(msg) + warnings.warn(msg, stacklevel=3) msg = "" if isinstance(from_schema.time_config, DatetimeRange) and time_zone_column is None: @@ -341,7 +335,7 @@ def _check_standard_time_zones(self) -> None: msg = "" time_zones = self._from_schema.time_config.time_zones for tz in time_zones: - if tz == "None": + if tz is None: msg += "\nChronify does not support None time zone in time_zone_column. " raise InvalidParameter(msg) if not is_standard_time_zone(tz): @@ -405,9 +399,10 @@ def generate_to_time_config(self) -> DatetimeRangeBase: raise InvalidParameter(msg) def generate_to_schema(self) -> TableSchema: - id_cols = self._from_schema.time_array_id_columns - if "time_zone" not in id_cols: - id_cols.append("time_zone") + id_cols = self._from_schema.time_array_id_columns.copy() + assert self.time_zone_column is not None # Validated in __init__ + if self.time_zone_column not in id_cols: + id_cols.append(self.time_zone_column) to_schema: TableSchema = self._from_schema.model_copy( update={ "name": f"{self._from_schema.name}_tz_converted", @@ -430,8 +425,7 @@ def localize_time_zone( mapping_schema, self._from_schema, self._to_schema, - self._engine, - self._metadata, + self._backend, TimeBasedDataAdjustment(), scratch_dir=scratch_dir, output_file=output_file, @@ -439,16 +433,14 @@ def localize_time_zone( ) def _get_time_zones(self) -> list[tzinfo | None]: - with self._engine.connect() as conn: - table = Table(self._from_schema.name, self._metadata) - stmt = ( - select(table.c[self.time_zone_column]) - .distinct() - .where(table.c[self.time_zone_column].is_not(None)) - ) - time_zones = read_database(stmt, conn, self._from_schema.time_config)[ - self.time_zone_column - ].to_list() + table = self._backend.table(self._from_schema.name) + query = ( + table.select(table[self.time_zone_column]) + .distinct() + .filter(table[self.time_zone_column].notnull()) + ) + df = read_query(self._backend, query, self._from_schema.time_config) + time_zones = df[self.time_zone_column].to_list() if "None" in time_zones and len(time_zones) > 1: msg = ( @@ -463,6 +455,9 @@ def _get_time_zones(self) -> list[tzinfo | None]: def _create_mapping(self) -> tuple[pd.DataFrame, MappingTableSchema]: """Create mapping dataframe for localizing tz-naive datetime to column time zones""" + assert isinstance( + self._from_schema.time_config, DatetimeRangeWithTZColumn + ) # Ensured by __init__ time_col = self._from_schema.time_config.time_column from_time_col = "from_" + time_col from_time_generator = make_time_range_generator(self._from_schema.time_config) @@ -484,6 +479,7 @@ def _create_mapping(self) -> tuple[pd.DataFrame, MappingTableSchema]: ) raise InvalidParameter(msg) + assert self.time_zone_column is not None # Validated in __init__ from_tz_col = "from_" + self.time_zone_column from_time_config = self._from_schema.time_config.model_copy( update={"time_column": from_time_col} diff --git a/src/chronify/utils/sqlalchemy_table.py b/src/chronify/utils/sqlalchemy_table.py deleted file mode 100644 index ec7e0b7..0000000 --- a/src/chronify/utils/sqlalchemy_table.py +++ /dev/null @@ -1,73 +0,0 @@ -# Copied this code from https://github.com/sqlalchemy/sqlalchemy/wiki/Views - -from typing import Any - -import sqlalchemy as sa -from sqlalchemy import Engine, MetaData, Selectable, TableClause -from sqlalchemy.ext import compiler -from sqlalchemy.schema import DDLElement -from sqlalchemy.sql import table -from sqlalchemy.sql.sqltypes import DATETIME, TIMESTAMP, TEXT - - -class CreateTable(DDLElement): - def __init__(self, name: str, selectable: Selectable) -> None: - self.name = name - self.selectable = selectable - - -class DropTable(DDLElement): - def __init__(self, name: str) -> None: - self.name = name - - -@compiler.compiles(CreateTable) -def _create_table(element: Any, compiler: Any, **kw: Any) -> str: - return "CREATE TABLE %s AS %s" % ( - element.name, - compiler.sql_compiler.process(element.selectable, literal_binds=True), - ) - - -@compiler.compiles(DropTable) -def _drop_table(element: Any, compiler: Any, **kw: Any) -> str: - return "DROP TABLE %s" % (element.name) - - -def _table_exists(ddl: Any, target: Any, connection: Any, **kw: Any) -> Any: - return ddl.name in sa.inspect(connection).get_table_names() - - -def _table_doesnt_exist(ddl: Any, target: Any, connection: Any, **kw: Any) -> bool: - return not _table_exists(ddl, target, connection, **kw) - - -def create_table( - name: str, selectable: Selectable, engine: Engine, metadata: MetaData -) -> TableClause: - """Create a table from a selectable.""" - table_ = table(name) - table_._columns._populate_separate_keys( - col._make_proxy(table_) - for col in selectable.selected_columns # type: ignore - ) - sa.event.listen( - metadata, - "after_create", - CreateTable(name, selectable).execute_if(callable_=_table_doesnt_exist), # type: ignore - ) - sa.event.listen(metadata, "before_drop", DropTable(name).execute_if(callable_=_table_exists)) # type: ignore - metadata.create_all(engine) - metadata.reflect(engine, views=True) - mtable = metadata.tables[name] - if engine.name == "sqlite": - # This is a workaround for a case we don't understand. - # In some cases the datetime column schema is set to NUMERIC when the real values - # are strings. - for col in table_._columns: - mcol = mtable.columns[col.name] - if ( - isinstance(col.type, TIMESTAMP) or isinstance(col.type, DATETIME) - ) and not isinstance(mcol.type, TEXT): - mcol.type = TEXT() - return table_ diff --git a/src/chronify/utils/sqlalchemy_view.py b/src/chronify/utils/sqlalchemy_view.py deleted file mode 100644 index bf4c495..0000000 --- a/src/chronify/utils/sqlalchemy_view.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copied this code from https://github.com/sqlalchemy/sqlalchemy/wiki/Views - -from typing import Any - -import sqlalchemy as sa -from sqlalchemy import Engine, MetaData, Selectable, TableClause -from sqlalchemy.ext import compiler -from sqlalchemy.schema import DDLElement -from sqlalchemy.sql import table - - -class CreateView(DDLElement): - def __init__(self, name: str, selectable: Selectable) -> None: - self.name = name - self.selectable = selectable - - -class DropView(DDLElement): - def __init__(self, name: str) -> None: - self.name = name - - -@compiler.compiles(CreateView) -def _create_view(element: Any, compiler: Any, **kw: Any) -> str: - return "CREATE VIEW %s AS %s" % ( - element.name, - compiler.sql_compiler.process(element.selectable, literal_binds=True), - ) - - -@compiler.compiles(DropView) -def _drop_view(element: Any, compiler: Any, **kw: Any) -> str: - return "DROP VIEW %s" % (element.name) - - -def _view_exists(ddl: Any, target: Any, connection: Any, **kw: Any) -> Any: - return ddl.name in sa.inspect(connection).get_view_names() - - -def _view_doesnt_exist(ddl: Any, target: Any, connection: Any, **kw: Any) -> bool: - return not _view_exists(ddl, target, connection, **kw) - - -def create_view( - name: str, selectable: Selectable, engine: Engine, metadata: MetaData -) -> TableClause: - """Create a view from a selectable.""" - view = table(name) - view._columns._populate_separate_keys( - col._make_proxy(view) - for col in selectable.selected_columns # type: ignore - ) - sa.event.listen( - metadata, - "after_create", - CreateView(name, selectable).execute_if( - callable_=_view_doesnt_exist # type: ignore - ), - ) - sa.event.listen( - metadata, - "before_drop", - DropView(name).execute_if( - callable_=_view_exists # type: ignore - ), - ) - metadata.create_all(engine) - metadata.reflect(engine, views=True) - return view diff --git a/tests/conftest.py b/tests/conftest.py index 254385c..a3af738 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,94 +1,90 @@ -import os -from typing import Any, Generator +from typing import Generator from pathlib import Path from tempfile import NamedTemporaryFile + import numpy as np import pandas as pd import pytest +from pyspark.sql import SparkSession -from sqlalchemy import Engine, create_engine, text - +from chronify.ibis.backend import IbisBackend, make_backend from chronify.models import TableSchema from chronify.store import Store from chronify.time import RepresentativePeriodFormat from chronify.time_configs import RepresentativePeriodTimeNTZ, RepresentativePeriodTimeTZ -ENGINES: dict[str, dict[str, Any]] = { - "duckdb": {"url": "duckdb:///:memory:", "connect_args": {}, "kwargs": {}}, - "sqlite": {"url": "sqlite:///:memory:", "connect_args": {}, "kwargs": {}}, -} -HIVE_URL = os.getenv("CHRONIFY_HIVE_URL") -if HIVE_URL is not None: - ENGINES["hive"] = {"url": HIVE_URL, "connect_args": {}, "kwargs": {}} +BACKENDS = ("duckdb", "sqlite", "spark") @pytest.fixture -def create_duckdb_engine() -> Engine: - """Return a sqlalchemy engine for DuckDB.""" - return create_engine("duckdb:///:memory:") +def create_duckdb_backend() -> IbisBackend: + """Return an Ibis backend for DuckDB.""" + return make_backend("duckdb") -@pytest.fixture(params=[x for x in ENGINES.keys() if x != "hive"]) -def iter_engines(request) -> Generator[Engine, None, None]: - """Return an iterable of sqlalchemy in-memory engines to test.""" - engine = ENGINES[request.param] - yield create_engine(engine["url"], *engine["connect_args"], **engine["kwargs"]) +@pytest.fixture(params=[x for x in BACKENDS if x != "spark"]) +def iter_backends(request) -> Generator[IbisBackend, None, None]: + """Return an iterable of Ibis backends to test.""" + yield make_backend(request.param) -@pytest.fixture(params=[x for x in ENGINES.keys() if x != "hive"]) -def iter_stores_by_engine(request) -> Generator[Store, None, None]: - """Return an iterable of stores with different engines to test. - Will only return engines that support data ingestion. +@pytest.fixture(params=[x for x in BACKENDS if x != "spark"]) +def iter_stores_by_backend(request) -> Generator[Store, None, None]: + """Return an iterable of stores with different backends to test. + Will only return backends that support data ingestion. """ - engine = ENGINES[request.param] - engine = create_engine(engine["url"], *engine["connect_args"], **engine["kwargs"]) - store = Store(engine=engine) + backend = request.param + store = Store(backend_name=backend) yield store store.dispose() -@pytest.fixture(params=ENGINES.keys()) -def iter_stores_by_engine_no_data_ingestion(request) -> Generator[Store, None, None]: - """Return an iterable of stores with different engines to test.""" - engine = ENGINES[request.param] - if engine["url"].startswith("hive"): - store = Store.create_new_hive_store( - engine["url"], *engine["connect_args"], drop_schema=True, **engine["kwargs"] +@pytest.fixture(params=BACKENDS) +def iter_stores_by_backend_no_data_ingestion(request, tmp_path) -> Generator[Store, None, None]: + """Return an iterable of stores with different backends to test.""" + backend = request.param + orig_tables_and_views: set[str] | None = set() + if backend == "spark": + # Use a temp directory for the spark warehouse to avoid conflicts + warehouse_dir = tmp_path / "spark-warehouse" + spark = ( + SparkSession.builder.appName("chronify_test") + .config("spark.sql.warehouse.dir", str(warehouse_dir)) + .getOrCreate() ) - orig_tables_and_views = set() - with store.engine.begin() as conn: - for row in conn.execute(text("SHOW TABLES")).all(): - orig_tables_and_views.add(row[1]) + store = Store.create_new_spark_store(session=spark) + for name in store._backend.list_tables(): + orig_tables_and_views.add(name) else: - eng = create_engine(engine["url"], *engine["connect_args"], **engine["kwargs"]) - store = Store(engine=eng) + store = Store(backend_name=backend) orig_tables_and_views = None yield store - if engine["url"].startswith("hive"): - with store.engine.begin() as conn: - for row in conn.execute(text("SHOW VIEWS")).all(): - name = row[1] - if name not in orig_tables_and_views: - conn.execute(text(f"DROP VIEW {name}")) - for row in conn.execute(text("SHOW TABLES")).all(): - name = row[1] - if name not in orig_tables_and_views: - conn.execute(text(f"DROP TABLE {name}")) - - -@pytest.fixture(params=[x for x in ENGINES.keys() if x != "hive"]) -def iter_engines_file(request, tmp_path) -> Generator[Engine, None, None]: - """Return an iterable of sqlalchemy file-based engines to test.""" - engine = ENGINES[request.param] + if backend == "spark" and orig_tables_and_views is not None: + # Cleanup views and tables created during test + for name in store._backend.list_tables(): + if name not in orig_tables_and_views: + try: + store._backend.drop_view(name, if_exists=True) + except Exception: + pass + try: + store._backend.drop_table(name, if_exists=True) + except Exception: + pass + + +@pytest.fixture(params=[x for x in BACKENDS if x != "spark"]) +def iter_backends_file(request, tmp_path) -> Generator[IbisBackend, None, None]: + """Return an iterable of Ibis file-based backends to test.""" + backend = request.param file_path = tmp_path / "store.db" - url = engine["url"].replace(":memory:", str(file_path)) - yield create_engine(url, *engine["connect_args"], **engine["kwargs"]) + yield make_backend(backend, file_path=file_path) -@pytest.fixture(params=[x for x in ENGINES.keys() if x != "hive"]) -def iter_engine_names(request) -> Generator[str, None, None]: - """Return an iterable of engine names.""" +@pytest.fixture(params=[x for x in BACKENDS if x != "spark"]) +def iter_backend_names(request) -> Generator[str, None, None]: + """Return an iterable of backend names.""" yield request.param diff --git a/tests/test_checker_representative_time.py b/tests/test_checker_representative_time.py index 48e7c7c..6658aa5 100644 --- a/tests/test_checker_representative_time.py +++ b/tests/test_checker_representative_time.py @@ -1,73 +1,69 @@ -from sqlalchemy import Engine, MetaData, Table import pytest import pandas as pd -from chronify.sqlalchemy.functions import write_database +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import write_table from chronify.models import TableSchema from chronify.time_series_checker import check_timestamps from chronify.exceptions import InvalidTable def ingest_data_and_check( - engine: Engine, df: pd.DataFrame, schema: TableSchema, error: tuple[any, str] + backend: IbisBackend, df: pd.DataFrame, schema: TableSchema, error: tuple ) -> None: - metadata = MetaData() - with engine.begin() as conn: - write_database(df, conn, schema.name, [schema.time_config], if_table_exists="replace") - metadata.reflect(engine, views=True) + write_table(backend, df, schema.name, [schema.time_config], if_exists="replace") - with engine.connect() as conn: - table = Table(schema.name, metadata) - if error: - with pytest.raises(error[0], match=error[1]): - check_timestamps(conn, table, schema) - else: - check_timestamps(conn, table, schema) + table = backend.table(schema.name) + if error: + with pytest.raises(error[0], match=error[1]): + check_timestamps(backend, table, schema) + else: + check_timestamps(backend, table, schema) -def test_one_week_per_month_by_hour(iter_engines: Engine, one_week_per_month_by_hour_table): +def test_one_week_per_month_by_hour(iter_backends: IbisBackend, one_week_per_month_by_hour_table): df, _, schema = one_week_per_month_by_hour_table error = () - ingest_data_and_check(iter_engines, df, schema, error) + ingest_data_and_check(iter_backends, df, schema, error) def test_one_week_per_month_by_hour_missing_data( - iter_engines: Engine, one_week_per_month_by_hour_table + iter_backends: IbisBackend, one_week_per_month_by_hour_table ): df, _, schema = one_week_per_month_by_hour_table df2 = df.loc[df["hour"] != 0].copy() error = (InvalidTable, "Mismatch number of timestamps") - ingest_data_and_check(iter_engines, df2, schema, error) + ingest_data_and_check(iter_backends, df2, schema, error) -def test_consistent_time_nulls(iter_engines: Engine, one_week_per_month_by_hour_table): +def test_consistent_time_nulls(iter_backends: IbisBackend, one_week_per_month_by_hour_table): df, _, schema = one_week_per_month_by_hour_table df.loc[len(df)] = [4.0, None, None, None, None] error = () - ingest_data_and_check(iter_engines, df, schema, error) + ingest_data_and_check(iter_backends, df, schema, error) -def test_inconsistent_time_nulls(iter_engines: Engine, one_week_per_month_by_hour_table): +def test_inconsistent_time_nulls(iter_backends: IbisBackend, one_week_per_month_by_hour_table): df, _, schema = one_week_per_month_by_hour_table df.loc[len(df)] = [4.0, None, 1.0, 2.0, 0.345] error = (InvalidTable, "If any time columns have a NULL value for a row") - ingest_data_and_check(iter_engines, df, schema, error) + ingest_data_and_check(iter_backends, df, schema, error) def test_one_weekday_day_and_one_weekend_day_per_month_by_hour( - iter_engines: Engine, one_weekday_day_and_one_weekend_day_per_month_by_hour_table + iter_backends: IbisBackend, one_weekday_day_and_one_weekend_day_per_month_by_hour_table ): df, _, schema = one_weekday_day_and_one_weekend_day_per_month_by_hour_table error = () - ingest_data_and_check(iter_engines, df, schema, error) + ingest_data_and_check(iter_backends, df, schema, error) def test_one_weekday_day_and_one_weekend_day_per_month_by_hour_wrong_data( - iter_engines: Engine, one_weekday_day_and_one_weekend_day_per_month_by_hour_table + iter_backends: IbisBackend, one_weekday_day_and_one_weekend_day_per_month_by_hour_table ): df, _, schema = one_weekday_day_and_one_weekend_day_per_month_by_hour_table df3 = df.copy() df3.loc[df3["month"] == 12, "month"] = 0 error = (InvalidTable, "Actual timestamps do not match expected timestamps") - ingest_data_and_check(iter_engines, df3, schema, error) + ingest_data_and_check(iter_backends, df3, schema, error) diff --git a/tests/test_csv_parser.py b/tests/test_csv_parser.py index 6ce81a0..a5e5f3b 100644 --- a/tests/test_csv_parser.py +++ b/tests/test_csv_parser.py @@ -37,19 +37,19 @@ def time_series_NYMDPV(): return temp_csv_file(header + data) -def test_NMDH_parser(time_series_NMDH, iter_engines): - store = Store(iter_engines) +def test_NMDH_parser(time_series_NMDH, iter_backends): + store = Store(backend=iter_backends) parser = CsvTimeSeriesParser(store) parser.ingest_to_datetime(time_series_NMDH, "test_NMDH", 2023, 48) -def test_NYMDH_parser(time_series_NYMDH, iter_engines): - store = Store(iter_engines) +def test_NYMDH_parser(time_series_NYMDH, iter_backends): + store = Store(backend=iter_backends) parser = CsvTimeSeriesParser(store) parser.ingest_to_datetime(time_series_NYMDH, "test_NYMDH", 2025, 48) -def test_NYMDPV_parser(time_series_NYMDPV, iter_engines): - store = Store(iter_engines) +def test_NYMDPV_parser(time_series_NYMDPV, iter_backends): + store = Store(backend=iter_backends) parser = CsvTimeSeriesParser(store) parser.ingest_to_datetime(time_series_NYMDPV, "test_NYMDPV", 2025, 24) diff --git a/tests/test_mapper_column_representative_to_datetime.py b/tests/test_mapper_column_representative_to_datetime.py index 53a43db..b2f5cd9 100644 --- a/tests/test_mapper_column_representative_to_datetime.py +++ b/tests/test_mapper_column_representative_to_datetime.py @@ -1,31 +1,23 @@ from datetime import timedelta, datetime -from pathlib import Path import pandas as pd import numpy as np import pytest -from sqlalchemy import MetaData from chronify.time_configs import ( YearMonthDayHourTimeNTZ, MonthDayHourTimeNTZ, DatetimeRange, - TimeConfig, YearMonthDayPeriodTimeNTZ, ) from chronify.models import TableSchema, PivotedTableSchema from chronify.store import Store -from chronify.sqlalchemy.functions import write_database, read_database +from chronify.ibis.functions import read_query from chronify.time_series_mapper import map_time @pytest.fixture -def iter_store(iter_engines): - return Store(engine=iter_engines) - - -def ingest_csv(csv_file: Path, conn, name: str, time_configs: list[TimeConfig]): - data = pd.read_csv(csv_file) - write_database(data, conn, name, time_configs, if_table_exists="replace") +def iter_store(iter_stores_by_backend): + return iter_stores_by_backend def test_MDH_mapper(time_series_NMDH, iter_store: Store): @@ -54,9 +46,6 @@ def test_MDH_mapper(time_series_NMDH, iter_store: Store): data = pd.read_csv(time_series_NMDH) iter_store.ingest_pivoted_table(data, pivoted_input_schema, from_schema) - metadata = MetaData() - metadata.reflect(iter_store.engine, views=True) - to_schema = TableSchema( name="test_MDH_datetime", value_column="value", @@ -69,13 +58,13 @@ def test_MDH_mapper(time_series_NMDH, iter_store: Store): time_array_id_columns=["name"], ) - map_time(iter_store.engine, metadata, from_schema, to_schema, check_mapped_timestamps=True) + map_time(iter_store._backend, from_schema, to_schema, check_mapped_timestamps=True) - with iter_store.engine.connect() as conn: - mapped_table = read_database( - f"SELECT * FROM {to_schema.name}", conn, to_schema.time_config - ) - assert np.array_equal(mapped_table["value"].to_numpy(), np.arange(25, 73)) + table = iter_store._backend.table(to_schema.name) + mapped_table = read_query(iter_store._backend, table, to_schema.time_config).sort_values( + "timestamp" + ) + assert np.array_equal(mapped_table["value"].to_numpy(), np.arange(25, 73)) def test_YMDH_mapper(time_series_NYMDH, iter_store): @@ -106,9 +95,6 @@ def test_YMDH_mapper(time_series_NYMDH, iter_store): data = pd.read_csv(time_series_NYMDH) iter_store.ingest_pivoted_table(data, pivoted_input_schema, from_schema) - metadata = MetaData() - metadata.reflect(iter_store.engine, views=True) - to_schema = TableSchema( name="test_YMDH_datetime", value_column="value", @@ -121,13 +107,13 @@ def test_YMDH_mapper(time_series_NYMDH, iter_store): time_array_id_columns=["name"], ) - map_time(iter_store.engine, metadata, from_schema, to_schema, check_mapped_timestamps=True) + map_time(iter_store._backend, from_schema, to_schema, check_mapped_timestamps=True) - with iter_store.engine.connect() as conn: - mapped_table = read_database( - f"SELECT * FROM {to_schema.name}", conn, to_schema.time_config - ) - assert np.array_equal(mapped_table["value"].to_numpy(), np.arange(25, 73)) + table = iter_store._backend.table(to_schema.name) + mapped_table = read_query(iter_store._backend, table, to_schema.time_config).sort_values( + "timestamp" + ) + assert np.array_equal(mapped_table["value"].to_numpy(), np.arange(25, 73)) def test_NYMDPV_mapper(time_series_NYMDPV, iter_store: Store): @@ -148,9 +134,6 @@ def test_NYMDPV_mapper(time_series_NYMDPV, iter_store: Store): data = pd.read_csv(time_series_NYMDPV) iter_store.ingest_table(data, from_schema, skip_time_checks=True) - metadata = MetaData() - metadata.reflect(iter_store.engine, views=True) - to_schema = TableSchema( name="test_YMDH_datetime", value_column="value", @@ -163,18 +146,18 @@ def test_NYMDPV_mapper(time_series_NYMDPV, iter_store: Store): time_array_id_columns=["name"], ) - map_time(iter_store.engine, metadata, from_schema, to_schema, check_mapped_timestamps=True) - - with iter_store.engine.connect() as conn: - mapped_table = read_database( - f"SELECT * FROM {to_schema.name}", conn, to_schema.time_config - ).sort_values("timestamp") - values = np.concatenate( - [ - np.ones(5) * 100, - np.ones(7) * 200, - np.ones(12) * 300, - np.ones(24) * 400, - ] - ) - assert np.array_equal(mapped_table["value"].to_numpy(), values) + map_time(iter_store._backend, from_schema, to_schema, check_mapped_timestamps=True) + + table = iter_store._backend.table(to_schema.name) + mapped_table = read_query(iter_store._backend, table, to_schema.time_config).sort_values( + "timestamp" + ) + values = np.concatenate( + [ + np.ones(5) * 100, + np.ones(7) * 200, + np.ones(12) * 300, + np.ones(24) * 400, + ] + ) + assert np.array_equal(mapped_table["value"].to_numpy(), values) diff --git a/tests/test_mapper_datetime_to_datetime.py b/tests/test_mapper_datetime_to_datetime.py index 5d969c2..25c3d05 100644 --- a/tests/test_mapper_datetime_to_datetime.py +++ b/tests/test_mapper_datetime_to_datetime.py @@ -6,9 +6,9 @@ import numpy as np import pandas as pd -from sqlalchemy import Engine, MetaData -from chronify.sqlalchemy.functions import read_database, write_database +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import read_query, write_table from chronify.time_series_mapper import map_time from chronify.time_configs import DatetimeRange from chronify.models import TableSchema @@ -58,42 +58,36 @@ def get_datetime_schema( def ingest_data( - engine: Engine, + backend: IbisBackend, df: pd.DataFrame, schema: TableSchema, ) -> None: - metadata = MetaData() - with engine.begin() as conn: - write_database(df, conn, schema.name, [schema.time_config], if_table_exists="replace") - metadata.reflect(engine, views=True) + write_table(backend, df, schema.name, [schema.time_config], if_exists="replace") def run_test_with_error( - engine: Engine, + backend: IbisBackend, df: pd.DataFrame, from_schema: TableSchema, to_schema: TableSchema, error: tuple[Any, str], ) -> None: - metadata = MetaData() - ingest_data(engine, df, from_schema) + ingest_data(backend, df, from_schema) with pytest.raises(error[0], match=error[1]): - map_time(engine, metadata, from_schema, to_schema, check_mapped_timestamps=True) + map_time(backend, from_schema, to_schema, check_mapped_timestamps=True) def get_mapped_results( - engine: Engine, + backend: IbisBackend, df: pd.DataFrame, from_schema: TableSchema, to_schema: TableSchema, ) -> pd.DataFrame: - metadata = MetaData() - ingest_data(engine, df, from_schema) - map_time(engine, metadata, from_schema, to_schema, check_mapped_timestamps=True) + ingest_data(backend, df, from_schema) + map_time(backend, from_schema, to_schema, check_mapped_timestamps=True) - with engine.connect() as conn: - query = f"select * from {to_schema.name}" - queried = read_database(query, conn, to_schema.time_config) + table = backend.table(to_schema.name) + queried = read_query(backend, table, to_schema.time_config) queried = queried.sort_values(by=["id", "timestamp"]).reset_index(drop=True)[df.columns] return queried @@ -159,7 +153,7 @@ def test_roll_time_using_shift_and_wrap() -> None: @pytest.mark.parametrize("tzinfo", [ZoneInfo("US/Eastern"), None]) def test_time_interval_shift( - iter_engines: Engine, + iter_backends: IbisBackend, tzinfo: tzinfo | None, ) -> None: from_schema = get_datetime_schema( @@ -168,14 +162,14 @@ def test_time_interval_shift( df = generate_datetime_dataframe(from_schema) to_schema = get_datetime_schema(2020, tzinfo, TimeIntervalType.PERIOD_ENDING, "to_table") - queried = get_mapped_results(iter_engines, df, from_schema, to_schema) + queried = get_mapped_results(iter_backends, df, from_schema, to_schema) check_time_shift_timestamps(df, queried, to_schema.time_config) check_time_shift_values(df, queried, from_schema.time_config, to_schema.time_config) @pytest.mark.parametrize("tzinfo", [ZoneInfo("US/Eastern"), None]) def test_time_interval_shift_different_time_ranges( - iter_engines: Engine, + iter_backends: IbisBackend, tzinfo: tzinfo | None, ) -> None: from_schema = get_datetime_schema( @@ -185,7 +179,7 @@ def test_time_interval_shift_different_time_ranges( to_schema = get_datetime_schema(2020, tzinfo, TimeIntervalType.PERIOD_ENDING, "to_table") to_schema.time_config.start += to_schema.time_config.resolution - queried = get_mapped_results(iter_engines, df, from_schema, to_schema) + queried = get_mapped_results(iter_backends, df, from_schema, to_schema) check_time_shift_timestamps(df, queried, to_schema.time_config) assert df["value"].equals(queried["value"]) @@ -199,7 +193,7 @@ def test_time_interval_shift_different_time_ranges( ], ) def test_time_shift_different_timezones( - iter_engines: Engine, tzinfo_tuple: tuple[tzinfo | None] + iter_backends: IbisBackend, tzinfo_tuple: tuple[tzinfo | None] ) -> None: from_schema = get_datetime_schema( 2020, tzinfo_tuple[0], TimeIntervalType.PERIOD_BEGINNING, "from_table" @@ -209,54 +203,53 @@ def test_time_shift_different_timezones( 2020, tzinfo_tuple[1], TimeIntervalType.PERIOD_ENDING, "to_table" ) - queried = get_mapped_results(iter_engines, df, from_schema, to_schema) + queried = get_mapped_results(iter_backends, df, from_schema, to_schema) check_time_shift_timestamps(df, queried, to_schema.time_config) check_time_shift_values(df, queried, from_schema.time_config, to_schema.time_config) def test_instantaneous_interval_type( - iter_engines: Engine, + iter_backends: IbisBackend, ) -> None: from_schema = get_datetime_schema(2020, None, TimeIntervalType.PERIOD_BEGINNING, "from_table") df = generate_datetime_dataframe(from_schema) to_schema = get_datetime_schema(2020, None, TimeIntervalType.INSTANTANEOUS, "to_table") error = (ConflictingInputsError, "If instantaneous time interval is used") - run_test_with_error(iter_engines, df, from_schema, to_schema, error) + run_test_with_error(iter_backends, df, from_schema, to_schema, error) def test_schema_compatibility( - iter_engines: Engine, + iter_backends: IbisBackend, ) -> None: from_schema = get_datetime_schema(2020, None, TimeIntervalType.PERIOD_BEGINNING, "from_table") df = generate_datetime_dataframe(from_schema) to_schema = get_datetime_schema(2020, None, TimeIntervalType.PERIOD_ENDING, "to_table") to_schema.time_array_id_columns += ["extra_column"] error = (ConflictingInputsError, ".* cannot produce the columns") - run_test_with_error(iter_engines, df, from_schema, to_schema, error) + run_test_with_error(iter_backends, df, from_schema, to_schema, error) def test_measurement_type_consistency( - iter_engines: Engine, + iter_backends: IbisBackend, ) -> None: from_schema = get_datetime_schema(2020, None, TimeIntervalType.PERIOD_BEGINNING, "from_table") df = generate_datetime_dataframe(from_schema) to_schema = get_datetime_schema(2020, None, TimeIntervalType.PERIOD_ENDING, "to_table") to_schema.time_config.measurement_type = MeasurementType.MAX error = (ConflictingInputsError, "Inconsistent measurement_types") - run_test_with_error(iter_engines, df, from_schema, to_schema, error) + run_test_with_error(iter_backends, df, from_schema, to_schema, error) -def test_duplicated_configs_in_write_database( - iter_engines: Engine, +def test_duplicated_configs_in_write_table( + iter_backends: IbisBackend, ) -> None: schema = get_datetime_schema(2020, None, TimeIntervalType.PERIOD_BEGINNING, "from_table") df = generate_datetime_dataframe(schema) configs = [schema.time_config, schema.time_config] # Ingest - with iter_engines.connect() as conn: - if conn.engine.name == "sqlite": - with pytest.raises(InvalidParameter, match="More than one datetime config found"): - write_database(df, conn, schema.name, configs, if_table_exists="replace") - else: - write_database(df, conn, schema.name, configs, if_table_exists="replace") + if iter_backends.name == "sqlite": + with pytest.raises(InvalidParameter, match="More than one datetime config found"): + write_table(iter_backends, df, schema.name, configs, if_exists="replace") + else: + write_table(iter_backends, df, schema.name, configs, if_exists="replace") diff --git a/tests/test_mapper_index_time_to_datetime.py b/tests/test_mapper_index_time_to_datetime.py index df18a30..6615c91 100644 --- a/tests/test_mapper_index_time_to_datetime.py +++ b/tests/test_mapper_index_time_to_datetime.py @@ -1,11 +1,11 @@ import pandas as pd -from sqlalchemy import Engine, MetaData import pytest from datetime import timedelta from zoneinfo import ZoneInfo from typing import Any, Optional -from chronify.sqlalchemy.functions import read_database, write_database +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import read_query, write_table from chronify.time_series_mapper import map_time from chronify.time_configs import ( DatetimeRange, @@ -139,7 +139,7 @@ def data_for_unaligned_time_mapping( def run_test( - engine: Engine, + backend: IbisBackend, df: pd.DataFrame, from_schema: TableSchema, to_schema: TableSchema, @@ -148,19 +148,13 @@ def run_test( wrap_time_allowed: bool = False, ) -> None: # Ingest - metadata = MetaData() - with engine.begin() as conn: - write_database( - df, conn, from_schema.name, [from_schema.time_config], if_table_exists="replace" - ) - metadata.reflect(engine, views=True) + write_table(backend, df, from_schema.name, [from_schema.time_config], if_exists="replace") # Map if error: with pytest.raises(error[0], match=error[1]): map_time( - engine, - metadata, + backend, from_schema, to_schema, data_adjustment=data_adjustment, @@ -169,8 +163,7 @@ def run_test( ) else: map_time( - engine, - metadata, + backend, from_schema, to_schema, data_adjustment=data_adjustment, @@ -179,10 +172,9 @@ def run_test( ) -def get_output_table(engine: Engine, to_schema: TableSchema) -> pd.DataFrame: - with engine.connect() as conn: - query = f"select * from {to_schema.name}" - queried = read_database(query, conn, to_schema.time_config) +def get_output_table(backend: IbisBackend, to_schema: TableSchema) -> pd.DataFrame: + table = backend.table(to_schema.name) + queried = read_query(backend, table, to_schema.time_config) return queried @@ -190,22 +182,22 @@ def get_output_table(engine: Engine, to_schema: TableSchema) -> pd.DataFrame: @pytest.mark.parametrize("interval_shift", [False, True]) @pytest.mark.parametrize("dst_std_time", [False, True]) def test_simple_mapping( - iter_engines: Engine, src_tz_naive: bool, interval_shift: bool, dst_std_time: bool + iter_backends: IbisBackend, src_tz_naive: bool, interval_shift: bool, dst_std_time: bool ) -> None: src_df, src_schema, dst_schema = data_for_simple_mapping( tz_naive=src_tz_naive, interval_shift=interval_shift, standard_time=dst_std_time ) error = None - run_test(iter_engines, src_df, src_schema, dst_schema, error) + run_test(iter_backends, src_df, src_schema, dst_schema, error) - dfo = get_output_table(iter_engines, dst_schema) + dfo = get_output_table(iter_backends, dst_schema) assert sorted(dfo["value"]) == sorted(src_df["value"]) @pytest.mark.parametrize("interval_shift", [False, True]) @pytest.mark.parametrize("dst_std_time", [False, True]) def test_unaligned_time_mapping( - iter_engines: Engine, interval_shift: bool, dst_std_time: bool + iter_backends: IbisBackend, interval_shift: bool, dst_std_time: bool ) -> None: src_df, src_schema, dst_schema = data_for_unaligned_time_mapping( interval_shift=interval_shift, standard_time=dst_std_time @@ -213,7 +205,7 @@ def test_unaligned_time_mapping( error = None wrap_time_allowed = True run_test( - iter_engines, + iter_backends, src_df, src_schema, dst_schema, @@ -221,18 +213,18 @@ def test_unaligned_time_mapping( wrap_time_allowed=wrap_time_allowed, ) - dfo = get_output_table(iter_engines, dst_schema) + dfo = get_output_table(iter_backends, dst_schema) assert sorted(dfo["value"]) == sorted(src_df["value"]) -def test_unaligned_time_mapping_without_wrap_time(iter_engines: Engine) -> None: +def test_unaligned_time_mapping_without_wrap_time(iter_backends: IbisBackend) -> None: src_df, src_schema, dst_schema = data_for_unaligned_time_mapping() error = ( ConflictingInputsError, "Length must match between", ) run_test( - iter_engines, + iter_backends, src_df, src_schema, dst_schema, @@ -244,7 +236,7 @@ def test_unaligned_time_mapping_without_wrap_time(iter_engines: Engine) -> None: @pytest.mark.parametrize("dst_std_time", [False, True]) @pytest.mark.parametrize("interpolate_fallback", [False, True]) def test_industrial_time_mapping( - iter_engines: Engine, + iter_backends: IbisBackend, interval_shift: bool, dst_std_time: bool, interpolate_fallback: bool, @@ -263,7 +255,7 @@ def test_industrial_time_mapping( daylight_saving_adjustment=DaylightSavingAdjustmentType.DROP_SPRING_FORWARD_DUPLICATE_FALLBACK ) run_test( - iter_engines, + iter_backends, src_df, src_schema, dst_schema, @@ -272,7 +264,7 @@ def test_industrial_time_mapping( wrap_time_allowed=True, ) - dfo = get_output_table(iter_engines, dst_schema) + dfo = get_output_table(iter_backends, dst_schema) dfo = dfo.sort_values(by=["time_zone", "value"]).reset_index(drop=True) # Check value associated with springforward hour is dropped @@ -306,7 +298,7 @@ def test_industrial_time_mapping( @pytest.mark.parametrize("dst_std_time", [False, True]) @pytest.mark.parametrize("interpolate_fallback", [False, True]) def test_industrial_time_subhourly( - iter_engines: Engine, + iter_backends: IbisBackend, dst_std_time: bool, interpolate_fallback: bool, ) -> None: @@ -323,7 +315,7 @@ def test_industrial_time_subhourly( daylight_saving_adjustment=DaylightSavingAdjustmentType.DROP_SPRING_FORWARD_DUPLICATE_FALLBACK ) run_test( - iter_engines, + iter_backends, src_df, src_schema, dst_schema, @@ -332,7 +324,7 @@ def test_industrial_time_subhourly( wrap_time_allowed=True, ) - dfo = get_output_table(iter_engines, dst_schema) + dfo = get_output_table(iter_backends, dst_schema) dfo = dfo.sort_values(by=["time_zone", "timestamp"]).reset_index(drop=True) dfo.loc[dfo["timestamp"].dt.date.astype(str) == "2018-11-04"] diff --git a/tests/test_mapper_representative_time_to_datetime.py b/tests/test_mapper_representative_time_to_datetime.py index 2574a8a..3bc658f 100644 --- a/tests/test_mapper_representative_time_to_datetime.py +++ b/tests/test_mapper_representative_time_to_datetime.py @@ -4,9 +4,9 @@ from typing import Any, Optional import pandas as pd -from sqlalchemy import Engine, MetaData -from chronify.sqlalchemy.functions import read_database, write_database +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import read_query, write_table from chronify.time_series_mapper import map_time from chronify.time_configs import DatetimeRange from chronify.models import TableSchema @@ -39,31 +39,25 @@ def get_datetime_schema(year: int, tzinfo: tzinfo | None) -> TableSchema: def run_test( - engine: Engine, + backend: IbisBackend, df: pd.DataFrame, from_schema: TableSchema, to_schema: TableSchema, error: Optional[tuple[Any, str]], ) -> None: # Ingest - metadata = MetaData() - with engine.begin() as conn: - write_database( - df, conn, from_schema.name, [from_schema.time_config], if_table_exists="replace" - ) - metadata.reflect(engine, views=True) + write_table(backend, df, from_schema.name, [from_schema.time_config], if_exists="replace") # Map if error: with pytest.raises(error[0], match=error[1]): - map_time(engine, metadata, from_schema, to_schema, check_mapped_timestamps=True) + map_time(backend, from_schema, to_schema, check_mapped_timestamps=True) else: - map_time(engine, metadata, from_schema, to_schema, check_mapped_timestamps=True) + map_time(backend, from_schema, to_schema, check_mapped_timestamps=True) # Check mapped table - with engine.connect() as conn: - query = f"select * from {to_schema.name}" - queried = read_database(query, conn, to_schema.time_config) + table = backend.table(to_schema.name) + queried = read_query(backend, table, to_schema.time_config) queried = queried.sort_values(by=["id", "timestamp"]).reset_index(drop=True) truth = generate_datetime_data(to_schema.time_config) @@ -113,7 +107,7 @@ def check_mapped_values(dfo: pd.DataFrame, dfi: pd.DataFrame, time_delta: timede @pytest.mark.parametrize("interval_shift", [False, True]) def test_one_week_per_month_by_hour_tz_naive( - iter_engines: Engine, + iter_backends: IbisBackend, one_week_per_month_by_hour_table: tuple[pd.DataFrame, int, TableSchema], interval_shift: bool, ) -> None: @@ -129,12 +123,12 @@ def test_one_week_per_month_by_hour_tz_naive( if interval_shift: to_schema.time_config.interval_type = TimeIntervalType.PERIOD_ENDING error = None - run_test(iter_engines, df, schema, to_schema, error) + run_test(iter_backends, df, schema, to_schema, error) @pytest.mark.parametrize("interval_shift", [False, True]) def test_one_week_per_month_by_hour_tz_aware( - iter_engines: Engine, + iter_backends: IbisBackend, one_week_per_month_by_hour_table_tz: tuple[pd.DataFrame, int, TableSchema], interval_shift: bool, ) -> None: @@ -145,12 +139,12 @@ def test_one_week_per_month_by_hour_tz_aware( if interval_shift: to_schema.time_config.interval_type = TimeIntervalType.PERIOD_ENDING error = None - run_test(iter_engines, df, schema, to_schema, error) + run_test(iter_backends, df, schema, to_schema, error) @pytest.mark.parametrize("interval_shift", [False, True]) def test_one_weekday_day_and_one_weekend_day_per_month_by_hour_tz_naive( - iter_engines: Engine, + iter_backends: IbisBackend, one_weekday_day_and_one_weekend_day_per_month_by_hour_table: tuple[ pd.DataFrame, int, TableSchema ], @@ -163,12 +157,12 @@ def test_one_weekday_day_and_one_weekend_day_per_month_by_hour_tz_naive( if interval_shift: to_schema.time_config.interval_type = TimeIntervalType.PERIOD_ENDING error = None - run_test(iter_engines, df, schema, to_schema, error) + run_test(iter_backends, df, schema, to_schema, error) @pytest.mark.parametrize("interval_shift", [False, True]) def test_one_weekday_day_and_one_weekend_day_per_month_by_hour_tz_aware( - iter_engines: Engine, + iter_backends: IbisBackend, one_weekday_day_and_one_weekend_day_per_month_by_hour_table_tz: tuple[ pd.DataFrame, int, TableSchema ], @@ -181,11 +175,11 @@ def test_one_weekday_day_and_one_weekend_day_per_month_by_hour_tz_aware( if interval_shift: to_schema.time_config.interval_type = TimeIntervalType.PERIOD_ENDING error = None - run_test(iter_engines, df, schema, to_schema, error) + run_test(iter_backends, df, schema, to_schema, error) def test_instantaneous_interval_type( - iter_engines: Engine, + iter_backends: IbisBackend, one_week_per_month_by_hour_table: tuple[pd.DataFrame, int, TableSchema], ) -> None: df, _, schema = one_week_per_month_by_hour_table @@ -193,4 +187,4 @@ def test_instantaneous_interval_type( to_schema = get_datetime_schema(2020, None) to_schema.time_config.interval_type = TimeIntervalType.INSTANTANEOUS error = None - run_test(iter_engines, df, schema, to_schema, error) + run_test(iter_backends, df, schema, to_schema, error) diff --git a/tests/test_models.py b/tests/test_models.py index 480931a..b149dfe 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,13 +1,13 @@ import pytest -from sqlalchemy import BigInteger, Boolean, DateTime, Double, Integer, String +import ibis.expr.datatypes as dt from chronify.models import ColumnDType, _check_name from chronify.exceptions import InvalidValue def test_column_dtypes() -> None: - ColumnDType(name="col1", dtype=Integer()) - for dtype in (BigInteger, Boolean, DateTime, Double, String): + ColumnDType(name="col1", dtype=dt.Int64()) + for dtype in (dt.Int64, dt.Boolean, dt.Timestamp, dt.Float64, dt.String): ColumnDType(name="col1", dtype=dtype()) for string_type in ("int", "bigint", "bool", "datetime", "float", "str"): diff --git a/tests/test_spark_backend.py b/tests/test_spark_backend.py new file mode 100644 index 0000000..f12a3fa --- /dev/null +++ b/tests/test_spark_backend.py @@ -0,0 +1,461 @@ +"""Tests for the Spark backend.""" + +from datetime import datetime, timedelta + +import numpy as np +import pandas as pd +import pytest + +from chronify import DatetimeRange, Store, TableSchema +from chronify.exceptions import InvalidTable + + +@pytest.fixture +def spark_store(tmp_path): + """Create a Spark store for testing.""" + from pyspark.sql import SparkSession + + # Use a temp directory for the spark warehouse to avoid conflicts + warehouse_dir = tmp_path / "spark-warehouse" + spark = ( + SparkSession.builder.appName("chronify_test") + .config("spark.sql.warehouse.dir", str(warehouse_dir)) + .getOrCreate() + ) + + store = Store.create_new_spark_store(session=spark) + orig_tables = set(store._backend.list_tables()) + yield store + # Cleanup views and tables created during test + for name in store._backend.list_tables(): + if name not in orig_tables: + try: + store._backend.drop_view(name, if_exists=True) + except Exception: + pass + try: + store._backend.drop_table(name, if_exists=True) + except Exception: + pass + + +@pytest.fixture +def sample_parquet_file(tmp_path): + """Create a sample Parquet file for testing.""" + initial_time = datetime(2020, 1, 1) + resolution = timedelta(hours=1) + length = 8784 # leap year + timestamps = pd.date_range(initial_time, periods=length, freq=resolution, unit="us") + + dfs = [] + for i in range(1, 4): + df = pd.DataFrame( + { + "timestamp": timestamps, + "id": i, + "value": np.random.random(length), + } + ) + dfs.append(df) + df = pd.concat(dfs, ignore_index=True) + parquet_path = tmp_path / "data.parquet" + df.to_parquet(parquet_path, index=False) + + schema = TableSchema( + name="devices", + value_column="value", + time_config=DatetimeRange( + time_column="timestamp", + start=initial_time, + length=length, + resolution=resolution, + ), + time_array_id_columns=["id"], + ) + return parquet_path, schema, df + + +class TestCreateViewFromParquet: + """Tests for Store.create_view_from_parquet with Spark backend.""" + + def test_create_view_from_parquet_basic(self, spark_store, sample_parquet_file): + """Test basic view creation from a Parquet file.""" + parquet_path, schema, original_df = sample_parquet_file + + spark_store.create_view_from_parquet(parquet_path, schema) + + assert spark_store.has_table(schema.name) + assert schema.name in spark_store.list_tables() + + df = spark_store.read_table(schema.name) + assert len(df) == len(original_df) + assert set(df.columns) == set(original_df.columns) + assert len(df["id"].unique()) == 3 + + def test_create_view_from_parquet_time_validation(self, spark_store, tmp_path): + """Test that time validation catches invalid data.""" + initial_time = datetime(2020, 1, 1) + resolution = timedelta(hours=1) + length = 100 + + # Create data with missing timestamps (only 90 rows instead of 100) + timestamps = pd.date_range(initial_time, periods=90, freq=resolution, unit="us") + df = pd.DataFrame( + { + "timestamp": timestamps, + "id": 1, + "value": np.random.random(90), + } + ) + parquet_path = tmp_path / "invalid_data.parquet" + df.to_parquet(parquet_path, index=False) + + schema = TableSchema( + name="invalid_devices", + value_column="value", + time_config=DatetimeRange( + time_column="timestamp", + start=initial_time, + length=length, + resolution=resolution, + ), + time_array_id_columns=["id"], + ) + + with pytest.raises(InvalidTable): + spark_store.create_view_from_parquet(parquet_path, schema) + + # View should not exist after validation failure + assert not spark_store.has_table(schema.name) + + def test_create_view_from_parquet_bypass_checks(self, spark_store, tmp_path): + """Test creating a view with bypass_checks=True.""" + initial_time = datetime(2020, 1, 1) + resolution = timedelta(hours=1) + length = 100 + + # Create incomplete data + timestamps = pd.date_range(initial_time, periods=50, freq=resolution, unit="us") + df = pd.DataFrame( + { + "timestamp": timestamps, + "id": 1, + "value": np.random.random(50), + } + ) + parquet_path = tmp_path / "partial_data.parquet" + df.to_parquet(parquet_path, index=False) + + schema = TableSchema( + name="partial_devices", + value_column="value", + time_config=DatetimeRange( + time_column="timestamp", + start=initial_time, + length=length, + resolution=resolution, + ), + time_array_id_columns=["id"], + ) + + # Should succeed with bypass_checks=True + spark_store.create_view_from_parquet(parquet_path, schema, bypass_checks=True) + assert spark_store.has_table(schema.name) + + def test_drop_view(self, spark_store, sample_parquet_file): + """Test dropping a view created from Parquet.""" + parquet_path, schema, _ = sample_parquet_file + + spark_store.create_view_from_parquet(parquet_path, schema) + assert spark_store.has_table(schema.name) + + spark_store.drop_view(schema.name) + assert not spark_store.has_table(schema.name) + + +class TestCreateView: + """Tests for Store.create_view with Spark backend.""" + + def test_create_view_from_existing_table(self, spark_store, sample_parquet_file): + """Test registering an Ibis table as a view from an existing persistent table.""" + parquet_path, base_schema, original_df = sample_parquet_file + + # First create a persistent view from parquet + spark_store.create_view_from_parquet(parquet_path, base_schema) + + # Now get the table and create another view from it + table = spark_store.get_table(base_schema.name) + + # Create a new schema for the derived view + schema = TableSchema( + name="ibis_view", + value_column="value", + time_config=base_schema.time_config, + time_array_id_columns=["id"], + ) + + spark_store.create_view(schema, table) + + assert spark_store.has_table(schema.name) + df = spark_store.read_table(schema.name) + assert len(df) == len(original_df) + + def test_create_view_with_filter(self, spark_store, sample_parquet_file): + """Test registering a filtered Ibis table as a view.""" + parquet_path, base_schema, _ = sample_parquet_file + + # First create a persistent view from parquet + spark_store.create_view_from_parquet(parquet_path, base_schema) + + # Get table and apply a filter + table = spark_store.get_table(base_schema.name) + filtered_table = table.filter(table["id"] == 1) + + # Schema for single id + schema = TableSchema( + name="filtered_view", + value_column="value", + time_config=base_schema.time_config, + time_array_id_columns=["id"], + ) + + spark_store.create_view(schema, filtered_table) + + df = spark_store.read_table(schema.name) + assert len(df["id"].unique()) == 1 + assert df["id"].iloc[0] == 1 + + def test_create_view_time_validation(self, spark_store, tmp_path): + """Test that create_view validates timestamps.""" + initial_time = datetime(2020, 1, 1) + resolution = timedelta(hours=1) + length = 100 + + # Create incomplete data + timestamps = pd.date_range(initial_time, periods=50, freq=resolution, unit="us") + df = pd.DataFrame( + { + "timestamp": timestamps, + "id": 1, + "value": np.random.random(50), + } + ) + parquet_path = tmp_path / "incomplete.parquet" + df.to_parquet(parquet_path, index=False) + + # First create a view from parquet with bypass_checks to get a persistent table + base_schema = TableSchema( + name="incomplete_base", + value_column="value", + time_config=DatetimeRange( + time_column="timestamp", + start=initial_time, + length=50, # Match actual data + resolution=resolution, + ), + time_array_id_columns=["id"], + ) + spark_store.create_view_from_parquet(parquet_path, base_schema) + + # Now try to create a view with stricter schema + table = spark_store.get_table(base_schema.name) + + schema = TableSchema( + name="incomplete_view", + value_column="value", + time_config=DatetimeRange( + time_column="timestamp", + start=initial_time, + length=length, # Expects 100 rows but only 50 exist + resolution=resolution, + ), + time_array_id_columns=["id"], + ) + + with pytest.raises(InvalidTable): + spark_store.create_view(schema, table) + + assert not spark_store.has_table(schema.name) + + def test_create_view_bypass_checks(self, spark_store, tmp_path): + """Test create_view with bypass_checks=True.""" + initial_time = datetime(2020, 1, 1) + resolution = timedelta(hours=1) + + # Create minimal data + timestamps = pd.date_range(initial_time, periods=10, freq=resolution, unit="us") + df = pd.DataFrame( + { + "timestamp": timestamps, + "id": 1, + "value": np.random.random(10), + } + ) + parquet_path = tmp_path / "minimal.parquet" + df.to_parquet(parquet_path, index=False) + + # First create a view from parquet with bypass_checks + base_schema = TableSchema( + name="minimal_base", + value_column="value", + time_config=DatetimeRange( + time_column="timestamp", + start=initial_time, + length=10, + resolution=resolution, + ), + time_array_id_columns=["id"], + ) + spark_store.create_view_from_parquet(parquet_path, base_schema) + + table = spark_store.get_table(base_schema.name) + + schema = TableSchema( + name="minimal_view", + value_column="value", + time_config=DatetimeRange( + time_column="timestamp", + start=initial_time, + length=1000, # Much larger than actual data + resolution=resolution, + ), + time_array_id_columns=["id"], + ) + + # Should succeed with bypass_checks=True + spark_store.create_view(schema, table, bypass_checks=True) + assert spark_store.has_table(schema.name) + + def test_create_view_with_transformation(self, spark_store, sample_parquet_file): + """Test create_view with a transformed Ibis table.""" + parquet_path, base_schema, _ = sample_parquet_file + + # First create a persistent view from parquet + spark_store.create_view_from_parquet(parquet_path, base_schema) + + table = spark_store.get_table(base_schema.name) + + # Apply transformation: multiply value by 2 + transformed = table.mutate(value=table["value"] * 2) + + schema = TableSchema( + name="transformed_view", + value_column="value", + time_config=base_schema.time_config, + time_array_id_columns=["id"], + ) + + spark_store.create_view(schema, transformed) + + df = spark_store.read_table(schema.name) + assert spark_store.has_table(schema.name) + assert len(df) == 8784 * 3 + + +class TestSparkStoreCreation: + """Tests for Store.create_new_spark_store.""" + + def test_create_new_spark_store(self): + """Test creating a new Spark store.""" + store = Store.create_new_spark_store() + assert store._backend.name == "spark" + + def test_create_spark_store_with_existing_session(self): + """Test creating a Spark store with an existing SparkSession.""" + from pyspark.sql import SparkSession + + spark = SparkSession.builder.appName("test_chronify").getOrCreate() + + store = Store.create_new_spark_store(session=spark) + assert store._backend.name == "spark" + assert store._backend._session is spark + + def test_spark_data_ingestion_not_supported(self, spark_store): + """Test that data ingestion raises NotImplementedError for Spark.""" + df = pd.DataFrame( + { + "timestamp": pd.date_range("2020-01-01", periods=10, freq="h"), + "id": 1, + "value": range(10), + } + ) + schema = TableSchema( + name="test_table", + value_column="value", + time_config=DatetimeRange( + time_column="timestamp", + start=datetime(2020, 1, 1), + length=10, + resolution=timedelta(hours=1), + ), + time_array_id_columns=["id"], + ) + + with pytest.raises(NotImplementedError, match="Data ingestion through Spark"): + spark_store.ingest_table(df, schema) + + +class TestSparkReadOperations: + """Tests for read operations with Spark backend.""" + + def test_read_table(self, spark_store, sample_parquet_file): + """Test reading a table.""" + parquet_path, schema, original_df = sample_parquet_file + spark_store.create_view_from_parquet(parquet_path, schema) + + df = spark_store.read_table(schema.name) + assert len(df) == len(original_df) + + def test_read_query(self, spark_store, sample_parquet_file): + """Test reading with a query.""" + parquet_path, schema, _ = sample_parquet_file + spark_store.create_view_from_parquet(parquet_path, schema) + + query = f"SELECT * FROM {schema.name} WHERE id = 1" + df = spark_store.read_query(schema.name, query) + assert len(df["id"].unique()) == 1 + assert df["id"].iloc[0] == 1 + + def test_read_query_with_ibis_expression(self, spark_store, sample_parquet_file): + """Test reading with an Ibis expression.""" + parquet_path, schema, _ = sample_parquet_file + spark_store.create_view_from_parquet(parquet_path, schema) + + table = spark_store.get_table(schema.name) + filtered = table.filter(table["id"] == 2) + + df = spark_store.read_query(schema.name, filtered) + assert len(df["id"].unique()) == 1 + assert df["id"].iloc[0] == 2 + + +class TestSparkParquetOutput: + """Tests for writing Parquet with Spark backend.""" + + def test_write_table_to_parquet(self, spark_store, sample_parquet_file, tmp_path): + """Test writing a table to Parquet.""" + parquet_path, schema, _ = sample_parquet_file + spark_store.create_view_from_parquet(parquet_path, schema) + + output_path = tmp_path / "output.parquet" + spark_store.write_table_to_parquet(schema.name, output_path) + + assert output_path.exists() + df = pd.read_parquet(output_path) + assert len(df) == 8784 * 3 + + def test_write_query_to_parquet(self, spark_store, sample_parquet_file, tmp_path): + """Test writing a query result to Parquet.""" + parquet_path, schema, _ = sample_parquet_file + spark_store.create_view_from_parquet(parquet_path, schema) + + table = spark_store.get_table(schema.name) + filtered = table.filter(table["id"] == 1) + + output_path = tmp_path / "filtered_output.parquet" + spark_store.write_query_to_parquet(filtered, output_path) + + assert output_path.exists() + df = pd.read_parquet(output_path) + assert len(df) == 8784 + assert df["id"].iloc[0] == 1 diff --git a/tests/test_store.py b/tests/test_store.py index 11a5ced..5f2362f 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -7,19 +7,10 @@ from itertools import chain import duckdb +import ibis.expr.datatypes as dt import numpy as np import pandas as pd import pytest -from sqlalchemy import ( - Connection, - DateTime, - Double, - Engine, - Integer, - Table, - create_engine, - select, -) from chronify.csv_io import read_csv from chronify.duckdb.functions import unpivot @@ -31,6 +22,8 @@ TableAlreadyExists, TableNotStored, ) +from chronify.ibis.backend import IbisBackend, make_backend +from chronify.ibis.functions import _convert_spark_output_for_datetime from chronify.models import ColumnDType, CsvTableSchema, PivotedTableSchema, TableSchema from chronify.store import Store from chronify.time import TimeIntervalType, DaylightSavingAdjustmentType @@ -61,10 +54,10 @@ def generators_schema(): src_schema = CsvTableSchema( time_config=time_config, column_dtypes=[ - ColumnDType(name="timestamp", dtype=DateTime(timezone=False)), - ColumnDType(name="gen1", dtype=Double()), - ColumnDType(name="gen2", dtype=Double()), - ColumnDType(name="gen3", dtype=Double()), + ColumnDType(name="timestamp", dtype=dt.Timestamp()), + ColumnDType(name="gen1", dtype=dt.Float64()), + ColumnDType(name="gen2", dtype=dt.Float64()), + ColumnDType(name="gen3", dtype=dt.Float64()), ], value_columns=["gen1", "gen2", "gen3"], pivoted_dimension_name="generator", @@ -107,11 +100,11 @@ def multiple_tables(): @pytest.mark.parametrize("use_time_zone", [True, False]) -def test_ingest_csv(iter_stores_by_engine: Store, tmp_path, generators_schema, use_time_zone): - store = iter_stores_by_engine +def test_ingest_csv(iter_stores_by_backend: Store, tmp_path, generators_schema, use_time_zone): + store = iter_stores_by_backend src_file, src_schema, dst_schema = generators_schema src_schema.column_dtypes[0] = ColumnDType( - name="timestamp", dtype=DateTime(timezone=use_time_zone) + name="timestamp", dtype=dt.Timestamp(timezone="UTC") if use_time_zone else dt.Timestamp() ) if use_time_zone: new_src_file = tmp_path / "gen_tz.csv" @@ -140,10 +133,13 @@ def test_ingest_csv(iter_stores_by_engine: Store, tmp_path, generators_schema, u src_schema2 = CsvTableSchema( time_config=src_schema.time_config, column_dtypes=[ - ColumnDType(name="timestamp", dtype=DateTime(timezone=use_time_zone)), - ColumnDType(name="g1b", dtype=Double()), - ColumnDType(name="g2b", dtype=Double()), - ColumnDType(name="g3b", dtype=Double()), + ColumnDType( + name="timestamp", + dtype=dt.Timestamp(timezone="UTC") if use_time_zone else dt.Timestamp(), + ), + ColumnDType(name="g1b", dtype=dt.Float64()), + ColumnDType(name="g2b", dtype=dt.Float64()), + ColumnDType(name="g3b", dtype=dt.Float64()), ], value_columns=["g1b", "g2b", "g3b"], pivoted_dimension_name="generator", @@ -168,11 +164,9 @@ def test_ingest_csv(iter_stores_by_engine: Store, tmp_path, generators_schema, u def test_ingest_csvs_with_rollback(tmp_path, multiple_tables): - # Python sqlite3 does not appear to support rollbacks with DDL statements. - # See discussion at https://bugs.python.org/issue10740. - # TODO: needs investigation - # Most users won't care...and will be using duckdb since it is the default. - store = Store(engine_name="duckdb") + # Note: Ibis doesn't support transactions in the same way as SQLAlchemy. + # This test verifies basic table creation behavior. + store = Store(backend_name="duckdb") tables, dst_schema = multiple_tables src_file1 = tmp_path / "file1.csv" src_file2 = tmp_path / "file2.csv" @@ -181,56 +175,35 @@ def test_ingest_csvs_with_rollback(tmp_path, multiple_tables): src_schema = CsvTableSchema( time_config=dst_schema.time_config, column_dtypes=[ - ColumnDType(name="timestamp", dtype=DateTime()), - ColumnDType(name="id", dtype=Integer()), - ColumnDType(name="value", dtype=Double()), + ColumnDType(name="timestamp", dtype=dt.Timestamp()), + ColumnDType(name="id", dtype=dt.Int64()), + ColumnDType(name="value", dtype=dt.Float64()), ], value_columns=[dst_schema.value_column], time_array_id_columns=dst_schema.time_array_id_columns, ) - def check_data(conn: Connection): - df = store.read_table(dst_schema.name, connection=conn) - assert len(df) == len(tables[0]) + len(tables[1]) - assert len(df.id.unique()) == 2 - - with store.engine.begin() as conn: - store.ingest_from_csvs((src_file1, src_file2), src_schema, dst_schema, connection=conn) - check_data(conn) - conn.rollback() - - store.update_metadata() - assert not store.has_table(dst_schema.name) - - with store.engine.begin() as conn: - store.ingest_from_csvs((src_file1, src_file2), src_schema, dst_schema, connection=conn) - check_data(conn) - - with store.engine.begin() as conn: - check_data(conn) + store.ingest_from_csvs((src_file1, src_file2), src_schema, dst_schema) + df = store.read_table(dst_schema.name) + assert len(df) == len(tables[0]) + len(tables[1]) + assert len(df.id.unique()) == 2 -@pytest.mark.parametrize("existing_connection", [False, True]) +@pytest.mark.parametrize("existing_connection", [False]) def test_ingest_multiple_tables( - iter_stores_by_engine: Store, multiple_tables, existing_connection: bool + iter_stores_by_backend: Store, multiple_tables, existing_connection: bool ): - store = iter_stores_by_engine + store = iter_stores_by_backend tables, schema = multiple_tables - if existing_connection: - store.ingest_tables(tables, schema) - else: - with store.engine.begin() as conn: - store.ingest_tables(tables, schema, connection=conn) - query = "SELECT * FROM devices WHERE id = ?" - params = (2,) - with store.engine.connect() as conn: - df = store.read_query("devices", query, params=params, connection=conn) + store.ingest_tables(tables, schema) + query = "SELECT * FROM devices WHERE id = 2" + df = store.read_query("devices", query) df["timestamp"] = df["timestamp"].astype("datetime64[ns]") assert df.equals(tables[1]) -def test_ingest_multiple_tables_error(iter_stores_by_engine: Store, multiple_tables): - store = iter_stores_by_engine +def test_ingest_multiple_tables_error(iter_stores_by_backend: Store, multiple_tables): + store = iter_stores_by_backend tables, schema = multiple_tables orig_value = tables[1].loc[8783]["id"] tables[1].loc[8783] = (tables[1].loc[8783]["timestamp"], 0.1, 99) @@ -240,28 +213,27 @@ def test_ingest_multiple_tables_error(iter_stores_by_engine: Store, multiple_tab tables[1].loc[8783] = (tables[1].loc[8783]["timestamp"], 0.1, orig_value) store.ingest_tables(tables, schema) - params = (2,) - df = store.read_query(schema.name, f"select * from {schema.name} where id=?", params=params) + df = store.read_query(schema.name, f"select * from {schema.name} where id=2") df["timestamp"] = df["timestamp"].astype("datetime64[ns]") assert df.equals(tables[1]) @pytest.mark.parametrize("use_pandas", [False, True]) -def test_ingest_pivoted_table(iter_stores_by_engine: Store, generators_schema, use_pandas: bool): - store = iter_stores_by_engine +def test_ingest_pivoted_table(iter_stores_by_backend: Store, generators_schema, use_pandas: bool): + store = iter_stores_by_backend src_file, src_schema, dst_schema = generators_schema pivoted_schema = PivotedTableSchema(**src_schema.model_dump(exclude={"column_dtypes"})) rel = read_csv(src_file, src_schema) input_table = rel.to_df() if use_pandas else rel store.ingest_pivoted_table(input_table, pivoted_schema, dst_schema) - table = store.get_table(dst_schema.name) - stmt = select(table).where(table.c.generator == "gen1") - df = store.read_query(dst_schema.name, stmt) + table = store._backend.table(dst_schema.name) + query = table.filter(table["generator"] == "gen1") + df = store._backend.execute(query) assert len(df) == 8784 -def test_ingest_invalid_csv(iter_stores_by_engine: Store, tmp_path, generators_schema): - store = iter_stores_by_engine +def test_ingest_invalid_csv(iter_stores_by_backend: Store, tmp_path, generators_schema): + store = iter_stores_by_backend src_file, src_schema, dst_schema = generators_schema lines = src_file.read_text().splitlines()[:-10] new_file = tmp_path / "data.csv" @@ -276,8 +248,8 @@ def test_ingest_invalid_csv(iter_stores_by_engine: Store, tmp_path, generators_s store.read_table(dst_schema.name) -def test_invalid_schema(iter_stores_by_engine: Store, generators_schema): - store = iter_stores_by_engine +def test_invalid_schema(iter_stores_by_backend: Store, generators_schema): + store = iter_stores_by_backend src_file, src_schema, dst_schema = generators_schema src_schema.value_columns = ["g1", "g2", "g3"] with pytest.raises(InvalidTable): @@ -285,9 +257,9 @@ def test_invalid_schema(iter_stores_by_engine: Store, generators_schema): def test_ingest_one_week_per_month_by_hour( - iter_stores_by_engine: Store, one_week_per_month_by_hour_table + iter_stores_by_backend: Store, one_week_per_month_by_hour_table ): - store = iter_stores_by_engine + store = iter_stores_by_backend df, num_time_arrays, schema = one_week_per_month_by_hour_table store.ingest_table(df, schema) @@ -300,9 +272,9 @@ def test_ingest_one_week_per_month_by_hour( def test_ingest_one_week_per_month_by_hour_invalid( - iter_stores_by_engine: Store, one_week_per_month_by_hour_table + iter_stores_by_backend: Store, one_week_per_month_by_hour_table ): - store = iter_stores_by_engine + store = iter_stores_by_backend df, _, schema = one_week_per_month_by_hour_table df_filtered = df[df["hour"] != 5] assert len(df_filtered) < len(df) @@ -311,9 +283,9 @@ def test_ingest_one_week_per_month_by_hour_invalid( store.ingest_table(df_filtered, schema) -def test_load_parquet(iter_stores_by_engine_no_data_ingestion: Store, tmp_path): - store = iter_stores_by_engine_no_data_ingestion - if store.engine.name == "sqlite": +def test_load_parquet(iter_stores_by_backend_no_data_ingestion: Store, tmp_path): + store = iter_stores_by_backend_no_data_ingestion + if store._backend.name == "sqlite": # SQLite doesn't support parquet return @@ -328,10 +300,10 @@ def test_load_parquet(iter_stores_by_engine_no_data_ingestion: Store, tmp_path): src_schema = CsvTableSchema( time_config=time_config, column_dtypes=[ - ColumnDType(name="timestamp", dtype=DateTime(timezone=False)), - ColumnDType(name="gen1", dtype=Double()), - ColumnDType(name="gen2", dtype=Double()), - ColumnDType(name="gen3", dtype=Double()), + ColumnDType(name="timestamp", dtype=dt.Timestamp()), + ColumnDType(name="gen1", dtype=dt.Float64()), + ColumnDType(name="gen2", dtype=dt.Float64()), + ColumnDType(name="gen3", dtype=dt.Float64()), ], value_columns=["gen1", "gen2", "gen3"], pivoted_dimension_name="generator", @@ -354,7 +326,7 @@ def test_load_parquet(iter_stores_by_engine_no_data_ingestion: Store, tmp_path): expected_timestamps = timestamp_generator.list_timestamps() all(df.timestamp.unique() == expected_timestamps) - # This adds test coverage for Hive. + # This adds test coverage for Spark. as_dict = dst_schema.model_dump() as_dict["name"] = "test_view" schema2 = TableSchema(**as_dict) @@ -383,12 +355,12 @@ def test_load_parquet(iter_stores_by_engine_no_data_ingestion: Store, tmp_path): ) def test_map_one_week_per_month_by_hour_to_datetime( tmp_path, - iter_stores_by_engine_no_data_ingestion: Store, + iter_stores_by_backend_no_data_ingestion: Store, one_week_per_month_by_hour_table, one_week_per_month_by_hour_table_tz, params: tuple[bool, int], ): - store = iter_stores_by_engine_no_data_ingestion + store = iter_stores_by_backend_no_data_ingestion use_time_zone, year = params if use_time_zone: df, num_time_arrays, src_schema = one_week_per_month_by_hour_table_tz @@ -407,7 +379,7 @@ def test_map_one_week_per_month_by_hour_to_datetime( ), time_array_id_columns=["id"], ) - if store.engine.name == "hive": + if store._backend.name == "spark": out_file = tmp_path / "data.parquet" df.to_parquet(out_file) store.create_view_from_parquet(out_file, src_schema) @@ -424,12 +396,8 @@ def test_map_one_week_per_month_by_hour_to_datetime( out_file = tmp_path / "out.parquet" assert not out_file.exists() - if store.engine.name == "sqlite": - with pytest.raises(NotImplementedError): - store.write_table_to_parquet(dst_schema.name, out_file) - else: - store.write_table_to_parquet(dst_schema.name, out_file, overwrite=True) - assert out_file.exists() + store.write_table_to_parquet(dst_schema.name, out_file, overwrite=True) + assert out_file.exists() with pytest.raises(TableAlreadyExists): store.map_table_time_config(src_schema.name, dst_schema, check_mapped_timestamps=True) @@ -437,9 +405,9 @@ def test_map_one_week_per_month_by_hour_to_datetime( @pytest.mark.parametrize("tzinfo", [ZoneInfo("EST"), None]) def test_map_datetime_to_datetime( - tmp_path, iter_stores_by_engine_no_data_ingestion: Store, tzinfo + tmp_path, iter_stores_by_backend_no_data_ingestion: Store, tzinfo ): - store = iter_stores_by_engine_no_data_ingestion + store = iter_stores_by_backend_no_data_ingestion time_array_len = 8784 year = 2020 @@ -461,10 +429,10 @@ def test_map_datetime_to_datetime( src_csv_schema = CsvTableSchema( time_config=src_time_config, column_dtypes=[ - ColumnDType(name="timestamp", dtype=DateTime(timezone=False)), - ColumnDType(name="gen1", dtype=Double()), - ColumnDType(name="gen2", dtype=Double()), - ColumnDType(name="gen3", dtype=Double()), + ColumnDType(name="timestamp", dtype=dt.Timestamp()), + ColumnDType(name="gen1", dtype=dt.Float64()), + ColumnDType(name="gen2", dtype=dt.Float64()), + ColumnDType(name="gen3", dtype=dt.Float64()), ], value_columns=["gen1", "gen2", "gen3"], pivoted_dimension_name="generator", @@ -479,30 +447,61 @@ def test_map_datetime_to_datetime( rel = read_csv(GENERATOR_TIME_SERIES_FILE, src_csv_schema) rel2 = unpivot(rel, ("gen1", "gen2", "gen3"), "generator", "value") # noqa: F841 + try: + df_debug = rel2.to_df() + print("DEBUG: Input DF dtypes:") + print(df_debug.dtypes) + print("DEBUG: Sample timestamp:") + print(df_debug["timestamp"].iloc[0]) + print(f"DEBUG: Sample timestamp tz: {df_debug['timestamp'].iloc[0].tzinfo}") + + target = pd.Timestamp("2020-11-01 08:00:00", tz="EST") + print(f"DEBUG: Looking for {target}") + subset = df_debug[df_debug["timestamp"] == target] + print(f"DEBUG: Found {len(subset)} rows.") + if not subset.empty: + print(subset.head()) + except Exception as e: + print(f"DEBUG: Error printing input DF: {e}") + src_schema = TableSchema( name="generators_pb", time_config=src_time_config, time_array_id_columns=["generator"], value_column="value", ) - if store.engine.name == "hive": + if store._backend.name == "spark": out_file = tmp_path / "data.parquet" rel2.to_df().to_parquet(out_file) store.create_view_from_parquet(out_file, src_schema) + + # DEBUG: Inspect Spark table + try: + print("DEBUG: Inspecting generators_pb in Spark") + spark = store._backend.con + df_spark = spark.sql( + "SELECT timestamp, CAST(timestamp as LONG) as ts_long FROM generators_pb WHERE timestamp >= '2020-11-01 12:00:00' AND timestamp <= '2020-11-01 14:00:00' ORDER BY timestamp" + ) + df_spark.show() + except Exception as e: + print(f"DEBUG: Error inspecting Spark table: {e}") + else: store.ingest_table(rel2, src_schema) - if tzinfo is None and store.engine.name != "sqlite": + if tzinfo is None and store._backend.name != "sqlite": output_file = tmp_path / "mapped_data" else: output_file = None store.map_table_time_config( src_schema.name, dst_schema, output_file=output_file, check_mapped_timestamps=True ) - if output_file is None or store.engine.name == "sqlite": + if output_file is None or store._backend.name == "sqlite": df2 = store.read_table(dst_schema.name) else: df2 = pd.read_parquet(output_file) + if store._backend.name == "spark": + _convert_spark_output_for_datetime(df2, dst_schema.time_config) assert len(df2) == time_array_len * 3 actual = sorted(df2["timestamp"].unique()) assert isinstance(src_schema.time_config, DatetimeRange) @@ -512,9 +511,9 @@ def test_map_datetime_to_datetime( def test_map_index_time_to_datetime( - tmp_path: Path, iter_stores_by_engine_no_data_ingestion: Store + tmp_path: Path, iter_stores_by_backend_no_data_ingestion: Store ) -> None: - store = iter_stores_by_engine_no_data_ingestion + store = iter_stores_by_backend_no_data_ingestion year = 2018 time_array_len = 8760 src_schema = TableSchema( @@ -557,14 +556,14 @@ def test_map_index_time_to_datetime( for i, time_zone in enumerate(time_zones) ] ) - if store.engine.name == "hive": + if store._backend.name == "spark": out_file = tmp_path / "data.parquet" src_df.to_parquet(out_file) store.create_view_from_parquet(out_file, src_schema) else: store.ingest_table(src_df, src_schema) - if store.engine.name != "sqlite": + if store._backend.name != "sqlite": output_file = tmp_path / "mapped_data" else: output_file = None @@ -578,7 +577,7 @@ def test_map_index_time_to_datetime( daylight_saving_adjustment=DaylightSavingAdjustmentType.DROP_SPRING_FORWARD_DUPLICATE_FALLBACK ), ) - if output_file is None or store.engine.name == "sqlite": + if output_file is None or store._backend.name == "sqlite": result = store.read_table(dst_schema.name) else: result = pd.read_parquet(output_file) @@ -611,77 +610,72 @@ def test_to_parquet(tmp_path, generators_schema): store = Store() store.ingest_from_csv(src_file, src_schema, dst_schema) filename = tmp_path / "data.parquet" - table = Table(dst_schema.name, store.metadata) - stmt = select(table).where(table.c.generator == "gen2") - store.write_query_to_parquet(stmt, filename, overwrite=True) + table = store._backend.table(dst_schema.name) + query = table.filter(table["generator"] == "gen2") + store.write_query_to_parquet(query, filename, overwrite=True) assert filename.exists() df = pd.read_parquet(filename) assert len(df) == 8784 -def test_load_existing_store(iter_engines_file, one_week_per_month_by_hour_table): - engine = iter_engines_file +def test_load_existing_store(iter_backends_file, one_week_per_month_by_hour_table): + backend = iter_backends_file df, _, schema = one_week_per_month_by_hour_table - store = Store(engine=engine) + store = Store(backend=backend) store.ingest_table(df, schema) df2 = store.read_table(schema.name) assert df2.equals(df) - file_path = Path(engine.url.database) + file_path = Path(backend._file_path) assert file_path.exists() - store2 = Store.load_from_file(engine_name=engine.name, file_path=file_path) + store2 = Store.load_from_file(backend_name=backend.name, file_path=file_path) df3 = store2.read_table(schema.name) assert df3.equals(df2) with pytest.raises(FileNotFoundError): - Store.load_from_file(engine_name=engine.name, file_path="./invalid/path") + Store.load_from_file(backend_name=backend.name, file_path="./invalid/path") -def test_create_methods(iter_engine_names, tmp_path): +def test_create_methods(iter_backend_names, tmp_path): path = tmp_path / "data.db" assert not path.exists() - Store.create_file_db(engine_name=iter_engine_names, file_path=path) + Store.create_file_db(backend_name=iter_backend_names, file_path=path) gc.collect() assert path.exists() with pytest.raises(InvalidOperation): - Store.create_file_db(engine_name=iter_engine_names, file_path=path) - Store.create_file_db(engine_name=iter_engine_names, file_path=path, overwrite=True) - Store.create_in_memory_db(engine_name=iter_engine_names) + Store.create_file_db(backend_name=iter_backend_names, file_path=path) + Store.create_file_db(backend_name=iter_backend_names, file_path=path, overwrite=True) + Store.create_in_memory_db(backend_name=iter_backend_names) -def test_invalid_hive_url(): - with pytest.raises(InvalidParameter): - Store.create_new_hive_store("duckdb:///:memory:") - - -def test_invalid_engine(): +def test_invalid_backend(): with pytest.raises(NotImplementedError): - Store(engine_name="hive") + Store(backend_name="invalid_backend") -def test_create_with_existing_engine(): - engine = create_engine("duckdb:///:memory:") - store = Store(engine=engine) - assert store.engine is engine +def test_create_with_existing_backend(): + backend = make_backend("duckdb") + store = Store(backend=backend) + assert store._backend is backend def test_create_with_sqlite(): - Store(engine_name="sqlite") + Store(backend_name="sqlite") def test_create_with_conflicting_parameters(): with pytest.raises(ConflictingInputsError): - Store(engine=create_engine("duckdb:///:memory:"), engine_name="duckdb") + Store(backend=make_backend("duckdb"), backend_name="duckdb") -def test_backup(iter_engines_file: Engine, one_week_per_month_by_hour_table, tmp_path): - engine = iter_engines_file +def test_backup(iter_backends_file: IbisBackend, one_week_per_month_by_hour_table, tmp_path): + backend = iter_backends_file df, _, schema = one_week_per_month_by_hour_table - store = Store(engine=engine) + store = Store(backend=backend) store.ingest_table(df, schema) dst_file = tmp_path / "backup.db" assert not dst_file.exists() store.backup(dst_file) assert dst_file.exists() - store2 = Store(engine_name=engine.name, file_path=dst_file) + store2 = Store(backend_name=backend.name, file_path=dst_file) df2 = store2.read_table(schema.name) assert df2.equals(df) @@ -697,9 +691,9 @@ def test_backup(iter_engines_file: Engine, one_week_per_month_by_hour_table, tmp def test_backup_not_allowed(one_week_per_month_by_hour_table, tmp_path): - engine = create_engine("duckdb:///:memory:") + backend = make_backend("duckdb") df, _, schema = one_week_per_month_by_hour_table - store = Store(engine=engine) + store = Store(backend=backend) store.ingest_table(df, schema) dst_file = tmp_path / "backup.db" assert not dst_file.exists() @@ -708,8 +702,8 @@ def test_backup_not_allowed(one_week_per_month_by_hour_table, tmp_path): assert not dst_file.exists() -def test_delete_rows(iter_stores_by_engine: Store, one_week_per_month_by_hour_table): - store = iter_stores_by_engine +def test_delete_rows(iter_stores_by_backend: Store, one_week_per_month_by_hour_table): + store = iter_stores_by_backend df, _, schema = one_week_per_month_by_hour_table store.ingest_table(df, schema) df2 = store.read_table(schema.name) @@ -720,8 +714,7 @@ def test_delete_rows(iter_stores_by_engine: Store, one_week_per_month_by_hour_ta store.delete_rows(schema.name, {"id": 2}) df3 = store.read_table(schema.name) assert sorted(df3["id"].unique()) == [1, 3] - with store.engine.begin() as conn: - store.delete_rows(schema.name, {"id": 1}, connection=conn) + store.delete_rows(schema.name, {"id": 1}) df4 = store.read_table(schema.name) assert sorted(df4["id"].unique()) == [3] store.delete_rows(schema.name, {"id": 3}) @@ -731,8 +724,8 @@ def test_delete_rows(iter_stores_by_engine: Store, one_week_per_month_by_hour_ta store.delete_rows(schema.name, {"id": 3}) -def test_drop_table(iter_stores_by_engine: Store, one_week_per_month_by_hour_table): - store = iter_stores_by_engine +def test_drop_table(iter_stores_by_backend: Store, one_week_per_month_by_hour_table): + store = iter_stores_by_backend df, _, schema = one_week_per_month_by_hour_table assert not store.list_tables() store.ingest_table(df, schema) @@ -746,23 +739,23 @@ def test_drop_table(iter_stores_by_engine: Store, one_week_per_month_by_hour_tab store.drop_table(schema.name) -def test_drop_view(iter_stores_by_engine: Store, one_week_per_month_by_hour_table): - store = iter_stores_by_engine +def test_drop_view(iter_stores_by_backend: Store, one_week_per_month_by_hour_table): + store = iter_stores_by_backend df, _, schema = one_week_per_month_by_hour_table store.ingest_table(df, schema) - table = Table(schema.name, store.metadata) - stmt = select(table).where(table.c.id == 1) + table = store._backend.table(schema.name) + query = table.filter(table["id"] == 1) inputs = schema.model_dump() inputs["name"] = make_temp_view_name() schema2 = TableSchema(**inputs) - store.create_view(schema2, stmt) + store.create_view(schema2, query) assert schema2.name in store.list_tables() store.drop_view(schema2.name) assert schema2.name not in store.list_tables() -def test_read_raw_query(iter_stores_by_engine: Store, one_week_per_month_by_hour_table): - store = iter_stores_by_engine +def test_read_raw_query(iter_stores_by_backend: Store, one_week_per_month_by_hour_table): + store = iter_stores_by_backend df, _, schema = one_week_per_month_by_hour_table store.ingest_table(df, schema) @@ -770,27 +763,19 @@ def test_read_raw_query(iter_stores_by_engine: Store, one_week_per_month_by_hour df2 = store.read_raw_query(query) assert df2.equals(df) - query = f"SELECT * FROM {schema.name} where id = ?" - params = (2,) - with store.engine.connect() as conn: - df2 = store.read_raw_query(query, params=params, connection=conn) - assert df2.equals(df[df["id"] == 2].reset_index(drop=True)) - -def test_check_timestamps(iter_stores_by_engine: Store, one_week_per_month_by_hour_table) -> None: - store = iter_stores_by_engine +def test_check_timestamps(iter_stores_by_backend: Store, one_week_per_month_by_hour_table) -> None: + store = iter_stores_by_backend df, _, schema = one_week_per_month_by_hour_table store.ingest_table(df, schema) store.check_timestamps(schema.name) - with store.engine.begin() as conn: - store.check_timestamps(schema.name, connection=conn) @pytest.mark.parametrize("to_time_zone", [ZoneInfo("US/Eastern"), ZoneInfo("US/Mountain"), None]) def test_convert_time_zone( - tmp_path, iter_stores_by_engine_no_data_ingestion: Store, to_time_zone: tzinfo | None + tmp_path, iter_stores_by_backend_no_data_ingestion: Store, to_time_zone: tzinfo | None ): - store = iter_stores_by_engine_no_data_ingestion + store = iter_stores_by_backend_no_data_ingestion time_array_len = 8784 year = 2020 tzinfo = ZoneInfo("EST") @@ -806,10 +791,10 @@ def test_convert_time_zone( src_csv_schema = CsvTableSchema( time_config=src_time_config, column_dtypes=[ - ColumnDType(name="timestamp", dtype=DateTime(timezone=False)), - ColumnDType(name="gen1", dtype=Double()), - ColumnDType(name="gen2", dtype=Double()), - ColumnDType(name="gen3", dtype=Double()), + ColumnDType(name="timestamp", dtype=dt.Timestamp()), + ColumnDType(name="gen1", dtype=dt.Float64()), + ColumnDType(name="gen2", dtype=dt.Float64()), + ColumnDType(name="gen3", dtype=dt.Float64()), ], value_columns=["gen1", "gen2", "gen3"], pivoted_dimension_name="generator", @@ -824,14 +809,14 @@ def test_convert_time_zone( time_array_id_columns=["generator"], value_column="value", ) - if store.engine.name == "hive": + if store._backend.name == "spark": out_file = tmp_path / "data.parquet" rel2.to_df().to_parquet(out_file) store.create_view_from_parquet(out_file, src_schema) else: store.ingest_table(rel2, src_schema) - if tzinfo is None and store.engine.name != "sqlite": + if tzinfo is None and store._backend.name != "sqlite": output_file = tmp_path / "mapped_data" else: output_file = None @@ -839,7 +824,7 @@ def test_convert_time_zone( dst_schema = store.convert_time_zone( src_schema.name, to_time_zone, output_file=output_file, check_mapped_timestamps=True ) - if output_file is None or store.engine.name == "sqlite": + if output_file is None or store._backend.name == "sqlite": df2 = store.read_table(dst_schema.name) else: df2 = pd.read_parquet(output_file) @@ -860,9 +845,9 @@ def test_convert_time_zone( @pytest.mark.parametrize("wrapped_time_allowed", [False, True]) def test_convert_time_zone_by_column( - tmp_path, iter_stores_by_engine_no_data_ingestion: Store, wrapped_time_allowed: bool + tmp_path, iter_stores_by_backend_no_data_ingestion: Store, wrapped_time_allowed: bool ): - store = iter_stores_by_engine_no_data_ingestion + store = iter_stores_by_backend_no_data_ingestion time_array_len = 8784 year = 2020 tzinfo = ZoneInfo("EST") @@ -878,10 +863,10 @@ def test_convert_time_zone_by_column( src_csv_schema = CsvTableSchema( time_config=src_time_config, column_dtypes=[ - ColumnDType(name="timestamp", dtype=DateTime(timezone=False)), - ColumnDType(name="gen1", dtype=Double()), - ColumnDType(name="gen2", dtype=Double()), - ColumnDType(name="gen3", dtype=Double()), + ColumnDType(name="timestamp", dtype=dt.Timestamp()), + ColumnDType(name="gen1", dtype=dt.Float64()), + ColumnDType(name="gen2", dtype=dt.Float64()), + ColumnDType(name="gen3", dtype=dt.Float64()), ], value_columns=["gen1", "gen2", "gen3"], pivoted_dimension_name="generator", @@ -901,14 +886,14 @@ def test_convert_time_zone_by_column( time_array_id_columns=["generator", "time_zone"], value_column="value", ) - if store.engine.name == "hive": + if store._backend.name == "spark": out_file = tmp_path / "data.parquet" rel2.to_df().to_parquet(out_file) store.create_view_from_parquet(out_file, src_schema) else: store.ingest_table(rel2, src_schema) - if tzinfo is None and store.engine.name != "sqlite": + if tzinfo is None and store._backend.name != "sqlite": output_file = tmp_path / "mapped_data" else: output_file = None @@ -920,7 +905,7 @@ def test_convert_time_zone_by_column( wrap_time_allowed=wrapped_time_allowed, check_mapped_timestamps=True, ) - if output_file is None or store.engine.name == "sqlite": + if output_file is None or store._backend.name == "sqlite": df2 = store.read_table(dst_schema.name) else: df2 = pd.read_parquet(output_file) @@ -943,10 +928,10 @@ def test_convert_time_zone_by_column( @pytest.mark.parametrize("to_time_zone", [ZoneInfo("UTC"), ZoneInfo("EST"), None]) def test_localize_time_zone( - tmp_path, iter_stores_by_engine_no_data_ingestion: Store, to_time_zone: tzinfo | None + tmp_path, iter_stores_by_backend_no_data_ingestion: Store, to_time_zone: tzinfo | None ): """Test time zone localization of tz-naive timestamps to a specified time zone.""" - store = iter_stores_by_engine_no_data_ingestion + store = iter_stores_by_backend_no_data_ingestion time_array_len = 8784 year = 2020 @@ -962,10 +947,10 @@ def test_localize_time_zone( src_csv_schema = CsvTableSchema( time_config=src_time_config, column_dtypes=[ - ColumnDType(name="timestamp", dtype=DateTime(timezone=False)), - ColumnDType(name="gen1", dtype=Double()), - ColumnDType(name="gen2", dtype=Double()), - ColumnDType(name="gen3", dtype=Double()), + ColumnDType(name="timestamp", dtype=dt.Timestamp()), + ColumnDType(name="gen1", dtype=dt.Float64()), + ColumnDType(name="gen2", dtype=dt.Float64()), + ColumnDType(name="gen3", dtype=dt.Float64()), ], value_columns=["gen1", "gen2", "gen3"], pivoted_dimension_name="generator", @@ -980,14 +965,14 @@ def test_localize_time_zone( time_array_id_columns=["generator"], value_column="value", ) - if store.engine.name == "hive": + if store._backend.name == "spark": out_file = tmp_path / "data.parquet" rel2.to_df().to_parquet(out_file) store.create_view_from_parquet(out_file, src_schema) else: store.ingest_table(rel2, src_schema) - if to_time_zone is None and store.engine.name != "sqlite": + if to_time_zone is None and store._backend.name != "sqlite": output_file = tmp_path / "mapped_data" else: output_file = None @@ -998,10 +983,12 @@ def test_localize_time_zone( output_file=output_file, check_mapped_timestamps=True, ) - if output_file is None or store.engine.name == "sqlite": + if output_file is None or store._backend.name == "sqlite": df2 = store.read_table(dst_schema.name) else: df2 = pd.read_parquet(output_file) + if store._backend.name == "spark": + _convert_spark_output_for_datetime(df2, dst_schema.time_config) df2["timestamp"] = pd.to_datetime(df2["timestamp"]) assert len(df2) == time_array_len * 3 actual = sorted(df2["timestamp"].unique()) @@ -1018,9 +1005,9 @@ def test_localize_time_zone( check_timestamp_lists(actual, expected) -def test_localize_time_zone_by_column(tmp_path, iter_stores_by_engine_no_data_ingestion: Store): +def test_localize_time_zone_by_column(tmp_path, iter_stores_by_backend_no_data_ingestion: Store): """Test time zone localization based on a time zone column.""" - store = iter_stores_by_engine_no_data_ingestion + store = iter_stores_by_backend_no_data_ingestion time_array_len = 8784 year = 2020 @@ -1036,10 +1023,10 @@ def test_localize_time_zone_by_column(tmp_path, iter_stores_by_engine_no_data_in src_csv_schema = CsvTableSchema( time_config=src_time_config, column_dtypes=[ - ColumnDType(name="timestamp", dtype=DateTime(timezone=False)), - ColumnDType(name="gen1", dtype=Double()), - ColumnDType(name="gen2", dtype=Double()), - ColumnDType(name="gen3", dtype=Double()), + ColumnDType(name="timestamp", dtype=dt.Timestamp()), + ColumnDType(name="gen1", dtype=dt.Float64()), + ColumnDType(name="gen2", dtype=dt.Float64()), + ColumnDType(name="gen3", dtype=dt.Float64()), ], value_columns=["gen1", "gen2", "gen3"], pivoted_dimension_name="generator", @@ -1059,14 +1046,14 @@ def test_localize_time_zone_by_column(tmp_path, iter_stores_by_engine_no_data_in time_array_id_columns=["generator", "time_zone"], value_column="value", ) - if store.engine.name == "hive": + if store._backend.name == "spark": out_file = tmp_path / "data.parquet" rel2.to_df().to_parquet(out_file) store.create_view_from_parquet(out_file, src_schema) else: store.ingest_table(rel2, src_schema) - if store.engine.name != "sqlite": + if store._backend.name != "sqlite": output_file = tmp_path / "mapped_data" else: output_file = None @@ -1077,7 +1064,7 @@ def test_localize_time_zone_by_column(tmp_path, iter_stores_by_engine_no_data_in output_file=output_file, check_mapped_timestamps=True, ) - if output_file is None or store.engine.name == "sqlite": + if output_file is None or store._backend.name == "sqlite": df2 = store.read_table(dst_schema.name) else: df2 = pd.read_parquet(output_file) diff --git a/tests/test_time_series_checker.py b/tests/test_time_series_checker.py index cf05c30..ea037b1 100644 --- a/tests/test_time_series_checker.py +++ b/tests/test_time_series_checker.py @@ -4,67 +4,63 @@ import pandas as pd import pytest -from sqlalchemy import ( - Engine, - MetaData, - Table, -) + +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import write_table from chronify.exceptions import InvalidTable from chronify.models import TableSchema -from chronify.sqlalchemy.functions import write_database from chronify.time import TimeIntervalType from chronify.time_configs import DatetimeRange from chronify.time_series_checker import check_timestamps -def test_valid_datetimes_with_tz(iter_engines: Engine) -> None: +def test_valid_datetimes_with_tz(iter_backends: IbisBackend) -> None: """Valid timestamps with time zones.""" - _run_test(iter_engines, *_get_inputs_for_valid_datetimes_with_tz()) + _run_test(iter_backends, *_get_inputs_for_valid_datetimes_with_tz()) -def test_valid_datetimes_without_tz(iter_engines: Engine) -> None: +def test_valid_datetimes_without_tz(iter_backends: IbisBackend) -> None: """Valid timestamps without time zones.""" - _run_test(iter_engines, *_get_inputs_for_valid_datetimes_without_tz()) + _run_test(iter_backends, *_get_inputs_for_valid_datetimes_without_tz()) -def test_invalid_datetimes(iter_engines: Engine) -> None: +def test_invalid_datetimes(iter_backends: IbisBackend) -> None: """Timestamps do not match the schema.""" - _run_test(iter_engines, *_get_inputs_for_incorrect_datetimes()) + _run_test(iter_backends, *_get_inputs_for_incorrect_datetimes()) -def test_invalid_datetime_length(iter_engines: Engine) -> None: +def test_invalid_datetime_length(iter_backends: IbisBackend) -> None: """Timestamps do not match the schema.""" - _run_test(iter_engines, *_get_inputs_for_incorrect_datetime_length()) + _run_test(iter_backends, *_get_inputs_for_incorrect_datetime_length()) -def test_mismatched_time_array_lengths(iter_engines: Engine) -> None: +def test_mismatched_time_array_lengths(iter_backends: IbisBackend) -> None: """Some time arrays have different lengths.""" - _run_test(iter_engines, *_get_inputs_for_mismatched_time_array_lengths()) + _run_test(iter_backends, *_get_inputs_for_mismatched_time_array_lengths()) -def test_incorrect_lengths(iter_engines: Engine) -> None: +def test_incorrect_lengths(iter_backends: IbisBackend) -> None: """All time arrays are consistent but have the wrong length.""" - _run_test(iter_engines, *_get_inputs_for_incorrect_lengths()) + _run_test(iter_backends, *_get_inputs_for_incorrect_lengths()) -def test_incorrect_time_arrays(iter_engines: Engine) -> None: +def test_incorrect_time_arrays(iter_backends: IbisBackend) -> None: """The time arrays form a complete set but are individually incorrect.""" - _run_test(iter_engines, *_get_inputs_for_incorrect_time_arrays()) + _run_test(iter_backends, *_get_inputs_for_incorrect_time_arrays()) -def test_incorrect_time_arrays_with_duplicates(iter_engines: Engine) -> None: +def test_incorrect_time_arrays_with_duplicates(iter_backends: IbisBackend) -> None: """The time arrays form a complete set but are individually incorrect.""" - _run_test(iter_engines, *_get_inputs_for_incorrect_time_arrays_with_duplicates()) + _run_test(iter_backends, *_get_inputs_for_incorrect_time_arrays_with_duplicates()) def _run_test( - engine: Engine, + backend: IbisBackend, df: pd.DataFrame, tzinfo: Optional[tzinfo], length: int, message: Optional[str], ) -> None: - metadata = MetaData() schema = TableSchema( name="generators", time_config=DatetimeRange( @@ -77,17 +73,14 @@ def _run_test( time_array_id_columns=["generator"], value_column="value", ) - with engine.begin() as conn: - write_database(df, conn, schema.name, [schema.time_config], if_table_exists="replace") - metadata.reflect(engine) - - with engine.connect() as conn: - table = Table(schema.name, metadata) - if message is None: - check_timestamps(conn, table, schema) - else: - with pytest.raises(InvalidTable, match=message): - check_timestamps(conn, table, schema) + write_table(backend, df, schema.name, [schema.time_config], if_exists="replace") + + table = backend.table(schema.name) + if message is None: + check_timestamps(backend, table, schema) + else: + with pytest.raises(InvalidTable, match=message): + check_timestamps(backend, table, schema) def _get_inputs_for_valid_datetimes_with_tz() -> tuple[pd.DataFrame, ZoneInfo, int, None]: diff --git a/tests/test_time_zone_converter.py b/tests/test_time_zone_converter.py index aca8cbf..27c7806 100644 --- a/tests/test_time_zone_converter.py +++ b/tests/test_time_zone_converter.py @@ -5,9 +5,9 @@ from typing import Any import pandas as pd -from sqlalchemy import Engine, MetaData -from chronify.sqlalchemy.functions import read_database, write_database +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import read_query, write_table from chronify.time_zone_converter import ( TimeZoneConverter, TimeZoneConverterByColumn, @@ -85,40 +85,33 @@ def get_datetime_schema( def ingest_data( - engine: Engine, - metadata: MetaData, + backend: IbisBackend, df: pd.DataFrame, schema: TableSchema, ) -> None: - with engine.begin() as conn: - write_database(df, conn, schema.name, [schema.time_config], if_table_exists="replace") - metadata.reflect(engine, views=True) + write_table(backend, df, schema.name, [schema.time_config], if_exists="replace") def get_mapped_dataframe( - engine: Engine, + backend: IbisBackend, table_name: str, time_config: DatetimeRange, ) -> pd.DataFrame: - with engine.connect() as conn: - query = f"select * from {table_name}" - queried = read_database(query, conn, time_config) + table = backend.table(table_name) + queried = read_query(backend, table, time_config) queried = queried.sort_values(by=["id", "timestamp"]).reset_index(drop=True) return queried def run_conversion( - engine: Engine, + backend: IbisBackend, df: pd.DataFrame, from_schema: TableSchema, to_time_zone: tzinfo | None, ) -> None: - metadata = MetaData() - ingest_data(engine, metadata, df, from_schema) - to_schema = convert_time_zone( - engine, metadata, from_schema, to_time_zone, check_mapped_timestamps=True - ) - dfo = get_mapped_dataframe(engine, to_schema.name, to_schema.time_config) + ingest_data(backend, df, from_schema) + to_schema = convert_time_zone(backend, from_schema, to_time_zone, check_mapped_timestamps=True) + dfo = get_mapped_dataframe(backend, to_schema.name, to_schema.time_config) assert df["value"].equals(dfo["value"]) if to_time_zone is None: expected = df["timestamp"].dt.tz_localize(None) @@ -128,22 +121,20 @@ def run_conversion( def run_conversion_to_column_time_zones( - engine: Engine, + backend: IbisBackend, df: pd.DataFrame, from_schema: TableSchema, wrap_time_allowed: bool, ) -> None: - metadata = MetaData() - ingest_data(engine, metadata, df, from_schema) + ingest_data(backend, df, from_schema) to_schema = convert_time_zone_by_column( - engine, - metadata, + backend, from_schema, "time_zone", wrap_time_allowed=wrap_time_allowed, check_mapped_timestamps=True, ) - dfo = get_mapped_dataframe(engine, to_schema.name, to_schema.time_config) + dfo = get_mapped_dataframe(backend, to_schema.name, to_schema.time_config) dfo = dfo[df.columns].sort_values(by="index").reset_index(drop=True) dfo["timestamp"] = pd.to_datetime(dfo["timestamp"]) # needed for engine 2, not sure why @@ -163,46 +154,45 @@ def run_conversion_to_column_time_zones( def run_conversion_with_error( - engine: Engine, + backend: IbisBackend, df: pd.DataFrame, from_schema: TableSchema, use_tz_col: bool, error: tuple[Any, str], ) -> None: - metadata = MetaData() - ingest_data(engine, metadata, df, from_schema) + ingest_data(backend, df, from_schema) with pytest.raises(error[0], match=error[1]): if use_tz_col: tzc = TimeZoneConverterByColumn( - engine, metadata, from_schema, "time_zone", wrap_time_allowed=False + backend, from_schema, "time_zone", wrap_time_allowed=False ) tzc.convert_time_zone(check_mapped_timestamps=True) else: - tzc2 = TimeZoneConverter(engine, metadata, from_schema, None) + tzc2 = TimeZoneConverter(backend, from_schema, None) tzc2.convert_time_zone(check_mapped_timestamps=True) -def test_src_table_no_time_zone(iter_engines: Engine) -> None: +def test_src_table_no_time_zone(iter_backends: IbisBackend) -> None: from_schema = get_datetime_schema(2018, None, TimeIntervalType.PERIOD_BEGINNING, "base_table") df = generate_datetime_dataframe(from_schema) error = (InvalidParameter, "Source schema time config start time must be timezone-aware") - run_conversion_with_error(iter_engines, df, from_schema, False, error) + run_conversion_with_error(iter_backends, df, from_schema, False, error) @pytest.mark.parametrize( "to_time_zone", [None, ZoneInfo("US/Central"), ZoneInfo("America/Los_Angeles")] ) -def test_time_conversion(iter_engines: Engine, to_time_zone: tzinfo | None) -> None: +def test_time_conversion(iter_backends: IbisBackend, to_time_zone: tzinfo | None) -> None: from_schema = get_datetime_schema( 2018, ZoneInfo("US/Mountain"), TimeIntervalType.PERIOD_BEGINNING, "base_table" ) df = generate_datetime_dataframe(from_schema) - run_conversion(iter_engines, df, from_schema, to_time_zone) + run_conversion(iter_backends, df, from_schema, to_time_zone) @pytest.mark.parametrize("wrap_time_allowed", [False, True]) def test_time_conversion_to_column_time_zones( - iter_engines: Engine, wrap_time_allowed: bool + iter_backends: IbisBackend, wrap_time_allowed: bool ) -> None: from_schema = get_datetime_schema( 2018, @@ -212,4 +202,4 @@ def test_time_conversion_to_column_time_zones( has_tz_col=True, ) df = generate_dataframe_with_tz_col(from_schema) - run_conversion_to_column_time_zones(iter_engines, df, from_schema, wrap_time_allowed) + run_conversion_to_column_time_zones(iter_backends, df, from_schema, wrap_time_allowed) diff --git a/tests/test_time_zone_localizer.py b/tests/test_time_zone_localizer.py index bc2968d..b149db0 100644 --- a/tests/test_time_zone_localizer.py +++ b/tests/test_time_zone_localizer.py @@ -5,9 +5,9 @@ from typing import Any import pandas as pd -from sqlalchemy import Engine, MetaData -from chronify.sqlalchemy.functions import read_database, write_database +from chronify.ibis.backend import IbisBackend +from chronify.ibis.functions import read_query, write_table from chronify.time_utils import get_standard_time_zone from chronify.time_zone_localizer import ( TimeZoneLocalizer, @@ -128,40 +128,35 @@ def get_datetime_with_tz_col_schema( def ingest_data( - engine: Engine, - metadata: MetaData, + backend: IbisBackend, df: pd.DataFrame, schema: TableSchema, ) -> None: - with engine.begin() as conn: - write_database(df, conn, schema.name, [schema.time_config], if_table_exists="replace") - metadata.reflect(engine, views=True) + write_table(backend, df, schema.name, [schema.time_config], if_exists="replace") def get_mapped_dataframe( - engine: Engine, + backend: IbisBackend, table_name: str, time_config: DatetimeRangeBase, ) -> pd.DataFrame: - with engine.connect() as conn: - query = f"select * from {table_name}" - queried = read_database(query, conn, time_config) + table = backend.table(table_name) + queried = read_query(backend, table, time_config) queried = queried.sort_values(by=["id", "timestamp"]).reset_index(drop=True) return queried def run_localization( - engine: Engine, + backend: IbisBackend, df: pd.DataFrame, from_schema: TableSchema, to_time_zone: tzinfo | None, ) -> None: - metadata = MetaData() - ingest_data(engine, metadata, df, from_schema) + ingest_data(backend, df, from_schema) to_schema = localize_time_zone( - engine, metadata, from_schema, to_time_zone, check_mapped_timestamps=True + backend, from_schema, to_time_zone, check_mapped_timestamps=True ) - dfo = get_mapped_dataframe(engine, to_schema.name, to_schema.time_config) + dfo = get_mapped_dataframe(backend, to_schema.name, to_schema.time_config) assert df["value"].equals(dfo["value"]) if to_time_zone is None: expected = df["timestamp"] @@ -173,19 +168,17 @@ def run_localization( def run_localization_to_column_time_zones( - engine: Engine, + backend: IbisBackend, df: pd.DataFrame, from_schema: TableSchema, ) -> None: - metadata = MetaData() - ingest_data(engine, metadata, df, from_schema) + ingest_data(backend, df, from_schema) to_schema = localize_time_zone_by_column( - engine, - metadata, + backend, from_schema, check_mapped_timestamps=True, ) - dfo = get_mapped_dataframe(engine, to_schema.name, to_schema.time_config) + dfo = get_mapped_dataframe(backend, to_schema.name, to_schema.time_config) dfo = dfo[df.columns].sort_values(by="index").reset_index(drop=True) dfo["timestamp"] = pd.to_datetime(dfo["timestamp"]) # needed for engine 2, not sure why assert df["value"].equals(dfo["value"]) @@ -201,48 +194,47 @@ def run_localization_to_column_time_zones( def run_localization_with_error( - engine: Engine, + backend: IbisBackend, df: pd.DataFrame, from_schema: TableSchema, error: tuple[Any, str], ) -> None: - metadata = MetaData() - ingest_data(engine, metadata, df, from_schema) + ingest_data(backend, df, from_schema) with pytest.raises(error[0], match=error[1]): - TimeZoneLocalizer(engine, metadata, from_schema, None).localize_time_zone( + TimeZoneLocalizer(backend, from_schema, None).localize_time_zone( check_mapped_timestamps=True ) def run_localization_by_column_with_error( - engine: Engine, + backend: IbisBackend, df: pd.DataFrame, from_schema: TableSchema, error: tuple[Any, str], time_zone_column: str | None = None, ) -> None: - metadata = MetaData() - ingest_data(engine, metadata, df, from_schema) + ingest_data(backend, df, from_schema) with pytest.raises(error[0], match=error[1]): TimeZoneLocalizerByColumn( - engine, - metadata, + backend, from_schema, time_zone_column=time_zone_column, ).localize_time_zone(check_mapped_timestamps=True) @pytest.mark.parametrize("to_time_zone", [None, ZoneInfo("EST")]) -def test_time_localization(iter_engines: Engine, to_time_zone: tzinfo | None) -> None: +def test_time_localization(iter_backends: IbisBackend, to_time_zone: tzinfo | None) -> None: from_schema = get_datetime_schema(2018, None, TimeIntervalType.PERIOD_BEGINNING, "base_table") df = generate_datetime_dataframe(from_schema) - run_localization(iter_engines, df, from_schema, to_time_zone) + run_localization(iter_backends, df, from_schema, to_time_zone) @pytest.mark.parametrize("from_time_tz", [None, ZoneInfo("US/Mountain"), ZoneInfo("MST")]) -def test_time_localization_by_column(iter_engines: Engine, from_time_tz: tzinfo | None) -> None: +def test_time_localization_by_column( + iter_backends: IbisBackend, from_time_tz: tzinfo | None +) -> None: from_schema = get_datetime_with_tz_col_schema( 2018, from_time_tz, @@ -251,36 +243,33 @@ def test_time_localization_by_column(iter_engines: Engine, from_time_tz: tzinfo standard_tz=True, ) df = generate_dataframe_with_tz_col(from_schema) - run_localization_to_column_time_zones(iter_engines, df, from_schema) + run_localization_to_column_time_zones(iter_backends, df, from_schema) # Error tests for TimeZoneLocalizer -def test_time_localizer_to_dst_time_error(iter_engines: Engine) -> None: +def test_time_localizer_to_dst_time_error(iter_backends: IbisBackend) -> None: """Test that TimeZoneLocalizer raises error when to_time_zone is a non standard time zone""" from_schema = get_datetime_schema(2018, None, TimeIntervalType.PERIOD_BEGINNING, "base_table") df = generate_datetime_dataframe(from_schema) to_time_zone = ZoneInfo("US/Mountain") # has DST - metadata = MetaData() - ingest_data(iter_engines, metadata, df, from_schema) + ingest_data(iter_backends, df, from_schema) with pytest.raises( InvalidParameter, match="TimeZoneLocalizer only supports standard time zones" ): - localize_time_zone( - iter_engines, metadata, from_schema, to_time_zone, check_mapped_timestamps=True - ) + localize_time_zone(iter_backends, from_schema, to_time_zone, check_mapped_timestamps=True) -def test_time_localizer_with_tz_aware_config_error(iter_engines: Engine) -> None: +def test_time_localizer_with_tz_aware_config_error(iter_backends: IbisBackend) -> None: """Test that TimeZoneLocalizer raises error when start time is tz-aware""" from_schema = get_datetime_schema( 2018, ZoneInfo("US/Mountain"), TimeIntervalType.PERIOD_BEGINNING, "base_table" ) df = generate_datetime_dataframe(from_schema) error = (InvalidParameter, "Source schema time config start time must be tz-naive") - run_localization_with_error(iter_engines, df, from_schema, error) + run_localization_with_error(iter_backends, df, from_schema, error) -def test_time_localizer_with_wrong_dtype_error(iter_engines: Engine) -> None: +def test_time_localizer_with_wrong_dtype_error(iter_backends: IbisBackend) -> None: """Test that TimeZoneLocalizer raises error when dtype is not TIMESTAMP_NTZ""" from_schema = get_datetime_schema(2018, None, TimeIntervalType.PERIOD_BEGINNING, "base_table") # Manually change dtype to TIMESTAMP_TZ to trigger error @@ -289,21 +278,21 @@ def test_time_localizer_with_wrong_dtype_error(iter_engines: Engine) -> None: ) df = generate_datetime_dataframe(from_schema) error = (InvalidParameter, "Source schema time config dtype must be TIMESTAMP_NTZ") - run_localization_with_error(iter_engines, df, from_schema, error) + run_localization_with_error(iter_backends, df, from_schema, error) -def test_time_localizer_with_datetime_range_with_tz_col_error(iter_engines: Engine) -> None: +def test_time_localizer_with_datetime_range_with_tz_col_error(iter_backends: IbisBackend) -> None: """Test that TimeZoneLocalizer raises error when time config is DatetimeRangeWithTZColumn""" from_schema = get_datetime_with_tz_col_schema( 2018, None, TimeIntervalType.PERIOD_BEGINNING, "base_table", standard_tz=True ) df = generate_dataframe_with_tz_col(from_schema) error = (InvalidParameter, "try using TimeZoneLocalizerByColumn") - run_localization_with_error(iter_engines, df, from_schema, error) + run_localization_with_error(iter_backends, df, from_schema, error) # Error tests for TimeZoneLocalizerByColumn -def test_time_localizer_by_column_to_dst_time_error(iter_engines: Engine) -> None: +def test_time_localizer_by_column_to_dst_time_error(iter_backends: IbisBackend) -> None: """Test that TimeZoneLocalizerByColumn raises error when to_time_zone is a non standard time zone""" from_schema = get_datetime_with_tz_col_schema( 2018, @@ -313,25 +302,22 @@ def test_time_localizer_by_column_to_dst_time_error(iter_engines: Engine) -> Non standard_tz=False, ) df = generate_dataframe_with_tz_col(from_schema) - metadata = MetaData() - ingest_data(iter_engines, metadata, df, from_schema) + ingest_data(iter_backends, df, from_schema) with pytest.raises( InvalidParameter, match="TimeZoneLocalizerByColumn only supports standard time zones" ): - localize_time_zone_by_column( - iter_engines, metadata, from_schema, check_mapped_timestamps=True - ) + localize_time_zone_by_column(iter_backends, from_schema, check_mapped_timestamps=True) -def test_time_localizer_by_column_missing_tz_column_error(iter_engines: Engine) -> None: +def test_time_localizer_by_column_missing_tz_column_error(iter_backends: IbisBackend) -> None: """Test that TimeZoneLocalizerByColumn raises error when time_zone_column is missing for DatetimeRange""" from_schema = get_datetime_schema(2018, None, TimeIntervalType.PERIOD_BEGINNING, "base_table") df = generate_datetime_dataframe(from_schema) error = (InvalidParameter, "time_zone_column must be provided") - run_localization_by_column_with_error(iter_engines, df, from_schema, error) + run_localization_by_column_with_error(iter_backends, df, from_schema, error) -def test_time_localizer_by_column_wrong_dtype_error(iter_engines: Engine) -> None: +def test_time_localizer_by_column_wrong_dtype_error(iter_backends: IbisBackend) -> None: """Test that TimeZoneLocalizerByColumn raises error when dtype is not TIMESTAMP_NTZ""" from_schema = get_datetime_with_tz_col_schema( 2018, None, TimeIntervalType.PERIOD_BEGINNING, "base_table", standard_tz=True @@ -342,20 +328,20 @@ def test_time_localizer_by_column_wrong_dtype_error(iter_engines: Engine) -> Non ) df = generate_dataframe_with_tz_col(from_schema) error = (InvalidParameter, "Source schema time config dtype must be TIMESTAMP_NTZ") - run_localization_by_column_with_error(iter_engines, df, from_schema, error) + run_localization_by_column_with_error(iter_backends, df, from_schema, error) -def test_time_localizer_by_column_non_standard_tz_error(iter_engines: Engine) -> None: +def test_time_localizer_by_column_non_standard_tz_error(iter_backends: IbisBackend) -> None: """Test that TimeZoneLocalizerByColumn raises error when time zones are not standard""" from_schema = get_datetime_with_tz_col_schema( 2018, None, TimeIntervalType.PERIOD_BEGINNING, "base_table", standard_tz=False ) df = generate_dataframe_with_tz_col(from_schema) error = (InvalidParameter, "is not a standard time zone") - run_localization_by_column_with_error(iter_engines, df, from_schema, error) + run_localization_by_column_with_error(iter_backends, df, from_schema, error) -def test_localize_time_zone_by_column_missing_tz_column_error(iter_engines: Engine) -> None: +def test_localize_time_zone_by_column_missing_tz_column_error(iter_backends: IbisBackend) -> None: """Test that localize_time_zone_by_column raises error when time_zone_column is None for DatetimeRange""" from_schema = get_datetime_schema(2018, None, TimeIntervalType.PERIOD_BEGINNING, "base_table") df = generate_datetime_dataframe(from_schema) @@ -364,5 +350,5 @@ def test_localize_time_zone_by_column_missing_tz_column_error(iter_engines: Engi "time_zone_column must be provided when source schema time config is of type DatetimeRange", ) run_localization_by_column_with_error( - iter_engines, df, from_schema, error, time_zone_column=None + iter_backends, df, from_schema, error, time_zone_column=None )