├── tests ├── superset_requirements.txt ├── conftest.py ├── unit │ ├── conftest.py │ └── test_databend_dialect.py ├── integration │ ├── test_sqlalchemy_integration.py │ └── conftest.py ├── test_table_options.py ├── test_copy_into.py ├── test_merge.py └── test_sqlalchemy.py ├── Makefile ├── Pipfile ├── pyproject.toml ├── databend_sqlalchemy ├── __init__.py ├── errors.py ├── provision.py ├── types.py ├── requirements.py ├── connector.py └── dml.py ├── .github └── workflows │ ├── ci.yml │ └── test.yml ├── .gitignore ├── setup.cfg ├── README.rst ├── LICENSE └── Pipfile.lock /tests/superset_requirements.txt: -------------------------------------------------------------------------------- 1 | apache_superset>=2.0.0 2 | flask==2.0.3 3 | MarkupSafe==2.0.1 4 | Jinja2==3.0.3 5 | werkzeug==2.0.3 6 | WTForms==2.3.3 -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | unittest: 2 | python -m pytest -s tests/unit 3 | 4 | integration: 5 | python -m pytest -s tests/integration 6 | 7 | testsuite: 8 | python -m pytest -n4 9 | 10 | install: 11 | pip install -e ".[dev]" 12 | -------------------------------------------------------------------------------- /Pipfile: -------------------------------------------------------------------------------- 1 | [[source]] 2 | url = "https://pypi.org/simple" 3 | verify_ssl = true 4 | name = "pypi" 5 | 6 | [packages] 7 | databend-sqlalchemy = {file = ".", editable = true} 8 | 9 | [dev-packages] 10 | black = "*" 11 | flake8 = "*" 12 | pytest = "*" 13 | pytest-xdist = "*" 14 | 15 | 16 | [requires] 17 | python_version = "3.11" 18 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools>=42", "wheel"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | # configure isort to be compatible with black 6 | # source: https://black.readthedocs.io/en/stable/compatible_configs.html#configuration 7 | [tool.isort] 8 | multi_line_output = 3 9 | include_trailing_comma = true 10 | force_grid_wrap = 0 11 | use_parentheses = true 12 | ensure_newline_before_comments = true 13 | 14 | [tool.black] 15 | exclude = ''' 16 | /( 17 | \.git 18 | | \.venv 19 | | dist 20 | )/ 21 | ''' 22 | -------------------------------------------------------------------------------- /databend_sqlalchemy/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | 4 | VERSION = (0, 5, 1) 5 | __version__ = ".".join(str(x) for x in VERSION) 6 | 7 | 8 | from .dml import ( 9 | Merge, 10 | WhenMergeUnMatchedClause, 11 | WhenMergeMatchedDeleteClause, 12 | WhenMergeMatchedUpdateClause, 13 | CopyIntoTable, 14 | CopyIntoLocation, 15 | CopyIntoTableOptions, 16 | CopyIntoLocationOptions, 17 | CSVFormat, 18 | TSVFormat, 19 | NDJSONFormat, 20 | ParquetFormat, 21 | ORCFormat, 22 | AVROFormat, 23 | AmazonS3, 24 | AzureBlobStorage, 25 | GoogleCloudStorage, 26 | FileColumnClause, 27 | StageClause, 28 | Compression, 29 | ) 30 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | pull_request: 8 | branches: 9 | - main 10 | 11 | concurrency: 12 | group: ${{ github.workflow }}-${{ github.ref }}-${{ github.event_name }} 13 | cancel-in-progress: true 14 | 15 | 16 | jobs: 17 | publish: 18 | runs-on: ubuntu-latest 19 | permissions: 20 | contents: read 21 | id-token: write 22 | steps: 23 | - uses: actions/checkout@v4 24 | - name: Build package 25 | run: | 26 | python -m pip install build 27 | python -m build 28 | - uses: pypa/gh-action-pypi-publish@release/v1 29 | if: github.event_name == 'push' && github.ref == 'refs/heads/main' 30 | with: 31 | packages-dir: dist/ 32 | skip-existing: true 33 | -------------------------------------------------------------------------------- /databend_sqlalchemy/errors.py: -------------------------------------------------------------------------------- 1 | class Error(Exception): 2 | code = None 3 | 4 | def __init__(self, message=None): 5 | self.message = message 6 | super(Error, self).__init__(message) 7 | 8 | def __str__(self): 9 | message = " " + self.message if self.message is not None else "" 10 | return "Code: {}.{}".format(self.code, message) 11 | 12 | 13 | class ServerException(Error): 14 | def __init__(self, message, code=None): 15 | self.message = message 16 | self.code = code 17 | super(ServerException, self).__init__(message) 18 | 19 | def __str__(self): 20 | return "Code: {}\n{}".format(self.code, self.message) 21 | 22 | 23 | class NotSupportedError(Error): 24 | def __init__(self, message, code=None): 25 | self.message = message 26 | super(NotSupportedError, self).__init__(message) 27 | 28 | def __str__(self): 29 | return "{}".format(self.message) 30 | -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | from sqlalchemy.dialects import registry 2 | import pytest 3 | 4 | registry.register("databend.databend", "databend_sqlalchemy.databend_dialect", "DatabendDialect") 5 | registry.register("databend", "databend_sqlalchemy.databend_dialect", "DatabendDialect") 6 | 7 | pytest.register_assert_rewrite("sa.testing.assertions") 8 | 9 | from sqlalchemy.testing.plugin.pytestplugin import * 10 | 11 | from packaging import version 12 | import sqlalchemy 13 | if version.parse(sqlalchemy.__version__) >= version.parse('2.0.0'): 14 | from sqlalchemy import event, text 15 | from sqlalchemy import Engine 16 | 17 | @event.listens_for(Engine, "connect") 18 | def receive_engine_connect(conn, r): 19 | cur = conn.cursor() 20 | cur.execute('SET global format_null_as_str = 0') 21 | cur.execute('SET global enable_geo_create_table = 1') 22 | try: 23 | cur.execute("SET global sequence_step_size = 1") 24 | except: 25 | pass # don't fail when this setting doesn't exist 26 | cur.close() 27 | -------------------------------------------------------------------------------- /tests/unit/conftest.py: -------------------------------------------------------------------------------- 1 | from unittest import mock 2 | from databend_sqlalchemy import databend_dialect 3 | from pytest import fixture 4 | 5 | 6 | class MockDBApi: 7 | class DatabaseError: 8 | pass 9 | 10 | class Error: 11 | pass 12 | 13 | class IntegrityError: 14 | pass 15 | 16 | class NotSupportedError: 17 | pass 18 | 19 | class OperationalError: 20 | pass 21 | 22 | class ProgrammingError: 23 | pass 24 | 25 | paramstyle = "" 26 | 27 | def execute(): 28 | pass 29 | 30 | def executemany(): 31 | pass 32 | 33 | def connect(): 34 | pass 35 | 36 | 37 | class MockCursor: 38 | def execute(): 39 | pass 40 | 41 | def executemany(): 42 | pass 43 | 44 | def fetchall(): 45 | pass 46 | 47 | def close(): 48 | pass 49 | 50 | 51 | @fixture 52 | def dialect() -> databend_dialect.DatabendDialect: 53 | return databend_dialect.DatabendDialect() 54 | 55 | 56 | @fixture 57 | def connection() -> mock.Mock(spec=MockDBApi): 58 | return mock.Mock(spec=MockDBApi) 59 | 60 | 61 | @fixture 62 | def cursor() -> mock.Mock(spec=MockCursor): 63 | return mock.Mock(spec=MockCursor) 64 | -------------------------------------------------------------------------------- /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | name: CI TEST 2 | 3 | on: 4 | push: 5 | branches: [main] 6 | pull_request: 7 | branches: [main] 8 | 9 | jobs: 10 | test: 11 | runs-on: ubuntu-latest 12 | services: 13 | databend: 14 | image: datafuselabs/databend:nightly 15 | env: 16 | QUERY_DEFAULT_USER: databend 17 | QUERY_DEFAULT_PASSWORD: databend 18 | MINIO_ENABLED: true 19 | ports: 20 | - 8000:8000 21 | - 9000:9000 22 | steps: 23 | - name: Checkout 24 | uses: actions/checkout@v4 25 | 26 | - name: Setup Python 27 | uses: actions/setup-python@v5 28 | with: 29 | python-version: "3.11" 30 | 31 | - name: Pip Install 32 | run: | 33 | pip install pipenv 34 | pipenv install --dev --skip-lock 35 | 36 | - name: Verify Service Running 37 | run: | 38 | cid=$(docker ps -a | grep databend | cut -d' ' -f1) 39 | docker logs ${cid} 40 | curl -v http://localhost:8000/v1/health 41 | 42 | - name: SQLAlchemy Test Suite 43 | env: 44 | TEST_DATABEND_DSN: "databend://databend:databend@localhost:8000/default?sslmode=disable" 45 | run: | 46 | pipenv run pytest -s . 47 | -------------------------------------------------------------------------------- /databend_sqlalchemy/provision.py: -------------------------------------------------------------------------------- 1 | 2 | from sqlalchemy.testing.provision import create_db 3 | from sqlalchemy.testing.provision import drop_db 4 | from sqlalchemy.testing.provision import configure_follower, update_db_opts, temp_table_keyword_args 5 | 6 | 7 | @create_db.for_db("databend") 8 | def _databend_create_db(cfg, eng, ident): 9 | with eng.begin() as conn: 10 | try: 11 | _databend_drop_db(cfg, conn, ident) 12 | except Exception: 13 | pass 14 | 15 | with eng.begin() as conn: 16 | conn.exec_driver_sql( 17 | "CREATE DATABASE IF NOT EXISTS %s " % ident 18 | ) 19 | conn.exec_driver_sql( 20 | "CREATE DATABASE IF NOT EXISTS %s_test_schema" % ident 21 | ) 22 | conn.exec_driver_sql( 23 | "CREATE DATABASE IF NOT EXISTS %s_test_schema_2" % ident 24 | ) 25 | 26 | 27 | @drop_db.for_db("databend") 28 | def _databend_drop_db(cfg, eng, ident): 29 | with eng.begin() as conn: 30 | conn.exec_driver_sql("DROP DATABASE IF EXISTS %s_test_schema" % ident) 31 | conn.exec_driver_sql("DROP DATABASE IF EXISTS %s_test_schema_2" % ident) 32 | conn.exec_driver_sql("DROP DATABASE IF EXISTS %s" % ident) 33 | 34 | @temp_table_keyword_args.for_db("databend") 35 | def _databend_temp_table_keyword_args(cfg, eng): 36 | return {"prefixes": ["TEMPORARY"]} 37 | 38 | 39 | @configure_follower.for_db("databend") 40 | def _databend_configure_follower(config, ident): 41 | config.test_schema = "%s_test_schema" % ident 42 | config.test_schema_2 = "%s_test_schema_2" % ident 43 | 44 | # Uncomment to debug SQL Statements in tests 45 | # @update_db_opts.for_db("databend") 46 | # def _databend_update_db_opts(db_url, db_opts): 47 | # db_opts["echo"] = True 48 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .envrc 2 | # Byte-compiled / optimized / DLL files 3 | __pycache__/ 4 | *.py[cod] 5 | *$py.class 6 | 7 | # C extensions 8 | *.so 9 | 10 | # Distribution / packaging 11 | .Python 12 | env/ 13 | build/ 14 | develop-eggs/ 15 | dist/ 16 | downloads/ 17 | eggs/ 18 | .eggs/ 19 | lib/ 20 | lib64/ 21 | parts/ 22 | sdist/ 23 | var/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *,cover 47 | .hypothesis/ 48 | 49 | # Translations 50 | *.mo 51 | *.pot 52 | 53 | # Django stuff: 54 | *.log 55 | local_settings.py 56 | 57 | # Flask stuff: 58 | instance/ 59 | .webassets-cache 60 | 61 | # Scrapy stuff: 62 | .scrapy 63 | 64 | # Sphinx documentation 65 | docs/_build/ 66 | 67 | # PyBuilder 68 | target/ 69 | 70 | # IPython Notebook 71 | .ipynb_checkpoints 72 | 73 | # pyenv 74 | .python-version 75 | 76 | # celery beat schedule file 77 | celerybeat-schedule 78 | 79 | # dotenv 80 | .env 81 | 82 | # virtualenv 83 | venv/ 84 | ENV/ 85 | 86 | # Spyder project settings 87 | .spyderproject 88 | 89 | # Rope project settings 90 | .ropeproject 91 | 92 | # PyCharm project settings 93 | .idea/ 94 | 95 | /dist 96 | /CHANGELOG.md 97 | /script/build 98 | 99 | # VS Code 100 | .vscode 101 | 102 | # IntelliJ 103 | .idea 104 | 105 | # macOS 106 | .DS_Store 107 | 108 | # vim 109 | *.swp 110 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | name = databend_sqlalchemy 3 | version = attr: databend_sqlalchemy.VERSION 4 | description = Sqlalchemy adapter for Databend 5 | long_description = file: README.rst 6 | long_description_content_type = text/markdown 7 | url = https://github.com/databendlabs/databend-sqlalchemy 8 | author = Databend Cloud 9 | author_email = hanshanjie@databend.com 10 | license = Apache-2.0 11 | license_file = LICENSE 12 | classifiers = 13 | License :: OSI Approved :: Apache Software License 14 | License :: OSI Approved :: MIT License 15 | Operating System :: OS Independent 16 | Programming Language :: Python :: 3 17 | Programming Language :: Python :: 3 :: Only 18 | Programming Language :: Python :: 3.7 19 | Programming Language :: Python :: 3.8 20 | Programming Language :: Python :: 3.9 21 | Programming Language :: Python :: 3.10 22 | Programming Language :: Python :: 3.11 23 | Programming Language :: Python :: 3.12 24 | project_urls = 25 | Bug Tracker = https://github.com/databendlabs/databend-sqlalchemy 26 | 27 | [options] 28 | packages = find: 29 | install_requires = 30 | databend-driver>=0.26.2 31 | sqlalchemy>=1.4 32 | python_requires = >=3.8 33 | package_dir = 34 | = . 35 | 36 | [options.packages.find] 37 | where = . 38 | 39 | [options.entry_points] 40 | sqlalchemy.dialects = 41 | databend = databend_sqlalchemy.databend_dialect:DatabendDialect 42 | databend.databend = databend_sqlalchemy.databend_dialect:DatabendDialect 43 | 44 | [options.extras_require] 45 | dev = 46 | devtools==0.7.0 47 | mock==4.0.3 48 | pre-commit==2.15.0 49 | pytest==8.1.1 50 | pytest-cov==3.0.0 51 | pytest-xdist==3.5.0 52 | sqlalchemy-stubs==0.4 53 | 54 | superset = apache_superset>=1.4.1 55 | 56 | [flake8] 57 | max-line-length = 88 58 | per-file-ignores = __init__.py:F401 59 | ignore = E203, W503 60 | ban-relative-imports = True 61 | inline-quotes = " 62 | 63 | [tool:pytest] 64 | addopts= --tb native -n1 -v -r fxX --maxfail=25 -p no:warnings 65 | python_files=tests/*test_*.py 66 | # python_functions=test_round_trip_same_named_column 67 | # log_level=DEBUG 68 | 69 | [sqla_testing] 70 | requirement_cls=databend_sqlalchemy.requirements:Requirements 71 | profile_file = .profiles.txt 72 | 73 | [db] 74 | default=databend://databend:databend@localhost:8000/default?sslmode=disable 75 | -------------------------------------------------------------------------------- /tests/integration/test_sqlalchemy_integration.py: -------------------------------------------------------------------------------- 1 | from datetime import date, datetime 2 | from decimal import Decimal 3 | 4 | import pytest 5 | from sqlalchemy import create_engine, text 6 | from sqlalchemy.engine.base import Connection, Engine 7 | from sqlalchemy.exc import OperationalError 8 | 9 | 10 | class TestDatabendDialect: 11 | def test_set_params( 12 | self, username: str, password: str, database_name: str, host_port_name: str 13 | ): 14 | engine = create_engine( 15 | f"databend://{username}:{password}@{host_port_name}/{database_name}?sslmode=disable" 16 | ) 17 | connection = engine.connect() 18 | result = connection.execute(text("SELECT 1")) 19 | assert len(result.fetchall()) == 1 20 | engine.dispose() 21 | 22 | def test_data_write(self, connection: Connection, fact_table_name: str): 23 | connection.execute( 24 | text(f"INSERT INTO {fact_table_name}(idx, dummy) VALUES (1, 'some_text')") 25 | ) 26 | result = connection.execute( 27 | text(f"SELECT * FROM {fact_table_name} WHERE idx=1") 28 | ) 29 | assert result.fetchall() == [(1, "some_text")] 30 | result = connection.execute(text(f"SELECT * FROM {fact_table_name}")) 31 | assert len(result.fetchall()) == 1 32 | 33 | def test_databend_types(self, connection: Connection): 34 | result = connection.execute(text("SELECT to_date('1896-01-01')")) 35 | print(str(date(1896, 1, 1))) 36 | assert result.fetchall() == [("1896-01-01",)] 37 | 38 | def test_has_table( 39 | self, engine: Engine, connection: Connection, fact_table_name: str 40 | ): 41 | results = engine.dialect.has_table(connection, fact_table_name) 42 | assert results == 1 43 | 44 | def test_get_columns( 45 | self, engine: Engine, connection: Connection, fact_table_name: str 46 | ): 47 | results = engine.dialect.get_columns(connection, fact_table_name) 48 | assert len(results) > 0 49 | row = results[0] 50 | assert isinstance(row, dict) 51 | row_keys = list(row.keys()) 52 | row_values = list(row.values()) 53 | assert row_keys[0] == "name" 54 | assert row_keys[1] == "type" 55 | assert row_keys[2] == "nullable" 56 | assert row_keys[3] == "default" 57 | print(row_values) 58 | -------------------------------------------------------------------------------- /tests/integration/conftest.py: -------------------------------------------------------------------------------- 1 | from logging import getLogger 2 | from os import environ 3 | 4 | from pytest import fixture 5 | from sqlalchemy import create_engine, text 6 | from sqlalchemy.engine.base import Connection, Engine 7 | 8 | LOGGER = getLogger(__name__) 9 | 10 | DATABASE_NAME = "default" 11 | USERNAME = "databend" 12 | PASSWORD = "databend" 13 | HOST_PORT = "localhost:8000" 14 | 15 | 16 | def must_env(var_name: str) -> str: 17 | assert var_name in environ, f"Expected {var_name} to be provided in environment" 18 | LOGGER.info(f"{var_name}: {environ[var_name]}") 19 | return environ[var_name] 20 | 21 | 22 | @fixture(scope="session") 23 | def host_port_name() -> str: 24 | return HOST_PORT 25 | 26 | 27 | @fixture(scope="session") 28 | def database_name() -> str: 29 | return DATABASE_NAME 30 | 31 | 32 | @fixture(scope="session") 33 | def username() -> str: 34 | return USERNAME 35 | 36 | 37 | @fixture(scope="session") 38 | def password() -> str: 39 | return PASSWORD 40 | 41 | 42 | @fixture(scope="session") 43 | def engine( 44 | username: str, password: str, host_port_name: str, database_name: str 45 | ) -> Engine: 46 | return create_engine( 47 | f"databend://{username}:{password}@{host_port_name}/{database_name}?sslmode=disable" 48 | ) 49 | 50 | 51 | @fixture(scope="session") 52 | def connection(engine: Engine) -> Connection: 53 | with engine.connect() as c: 54 | yield c 55 | 56 | 57 | @fixture(scope="class") 58 | def fact_table_name() -> str: 59 | return "test_alchemy" 60 | 61 | 62 | @fixture(scope="class") 63 | def dimension_table_name() -> str: 64 | return "test_alchemy_dimension" 65 | 66 | 67 | @fixture(scope="class", autouse=True) 68 | def setup_test_tables( 69 | connection: Connection, 70 | engine: Engine, 71 | fact_table_name: str, 72 | dimension_table_name: str, 73 | ): 74 | connection.execute( 75 | text( 76 | f""" 77 | CREATE TABLE IF NOT EXISTS {fact_table_name} 78 | ( 79 | idx INT, 80 | dummy VARCHAR 81 | ); 82 | """ 83 | ) 84 | ) 85 | connection.execute( 86 | text( 87 | f""" 88 | CREATE TABLE IF NOT EXISTS {dimension_table_name} 89 | ( 90 | idx INT, 91 | dummy VARCHAR 92 | ); 93 | """ 94 | ) 95 | ) 96 | assert engine.dialect.has_table(connection, fact_table_name) 97 | assert engine.dialect.has_table(connection, dimension_table_name) 98 | yield 99 | # Teardown 100 | connection.execute(text(f"DROP TABLE IF EXISTS {fact_table_name};")) 101 | connection.execute(text(f"DROP TABLE IF EXISTS {dimension_table_name};")) 102 | assert not engine.dialect.has_table(connection, fact_table_name) 103 | assert not engine.dialect.has_table(connection, dimension_table_name) 104 | -------------------------------------------------------------------------------- /databend_sqlalchemy/types.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | import datetime as dt 4 | from typing import Optional, Type, Any 5 | 6 | from sqlalchemy import func 7 | from sqlalchemy.engine.interfaces import Dialect 8 | from sqlalchemy.sql import sqltypes 9 | from sqlalchemy.sql import type_api 10 | 11 | 12 | # ToDo - This is perhaps how numeric should be defined 13 | # class NUMERIC(sqltypes.Numeric): 14 | # def result_processor(self, dialect, type_): 15 | # 16 | # orig = super().result_processor(dialect, type_) 17 | # 18 | # def process(value): 19 | # if value is not None: 20 | # if self.decimal_return_scale: 21 | # value = decimal.Decimal(f'{value:.{self.decimal_return_scale}f}') 22 | # else: 23 | # value = decimal.Decimal(value) 24 | # if orig: 25 | # return orig(value) 26 | # return value 27 | # 28 | # return process 29 | 30 | 31 | class INTERVAL(type_api.NativeForEmulated, sqltypes._AbstractInterval): 32 | """Databend INTERVAL type.""" 33 | 34 | __visit_name__ = "INTERVAL" 35 | native = True 36 | 37 | def __init__( 38 | self, precision: Optional[int] = None, fields: Optional[str] = None 39 | ) -> None: 40 | """Construct an INTERVAL. 41 | 42 | :param precision: optional integer precision value 43 | :param fields: string fields specifier. allows storage of fields 44 | to be limited, such as ``"YEAR"``, ``"MONTH"``, ``"DAY TO HOUR"``, 45 | etc. 46 | 47 | """ 48 | self.precision = precision 49 | self.fields = fields 50 | 51 | @classmethod 52 | def adapt_emulated_to_native( 53 | cls, 54 | interval: sqltypes.Interval, 55 | **kw: Any, # type: ignore[override] 56 | ) -> INTERVAL: 57 | return INTERVAL(precision=interval.second_precision) 58 | 59 | @property 60 | def _type_affinity(self) -> Type[sqltypes.Interval]: 61 | return sqltypes.Interval 62 | 63 | def as_generic(self, allow_nulltype: bool = False) -> sqltypes.Interval: 64 | return sqltypes.Interval(native=True, second_precision=self.precision) 65 | 66 | @property 67 | def python_type(self) -> Type[dt.timedelta]: 68 | return dt.timedelta 69 | 70 | def literal_processor( 71 | self, dialect: Dialect 72 | ) -> Optional[type_api._LiteralProcessorType[dt.timedelta]]: 73 | def process(value: dt.timedelta) -> str: 74 | return f"to_interval('{value.total_seconds()} seconds')" 75 | 76 | return process 77 | 78 | 79 | class TINYINT(sqltypes.Integer): 80 | __visit_name__ = "TINYINT" 81 | native = True 82 | 83 | 84 | class DOUBLE(sqltypes.Float): 85 | __visit_name__ = "DOUBLE" 86 | native = True 87 | 88 | 89 | class FLOAT(sqltypes.Float): 90 | __visit_name__ = "FLOAT" 91 | native = True 92 | 93 | 94 | # The “CamelCase” types are to the greatest degree possible database agnostic 95 | 96 | # For these datatypes, specific SQLAlchemy dialects provide backend-specific “UPPERCASE” datatypes, for a SQL type that has no analogue on other backends 97 | 98 | 99 | class BITMAP(sqltypes.TypeEngine): 100 | __visit_name__ = "BITMAP" 101 | render_bind_cast = True 102 | 103 | def __init__(self, **kwargs): 104 | super(BITMAP, self).__init__() 105 | 106 | def process_result_value(self, value, dialect): 107 | if value is None: 108 | return None 109 | # Databend returns bitmaps as strings of comma-separated integers 110 | return set(int(x) for x in value.split(',') if x) 111 | 112 | def bind_expression(self, bindvalue): 113 | return func.to_bitmap(bindvalue, type_=self) 114 | 115 | def column_expression(self, col): 116 | # Convert bitmap to string using a custom function 117 | return func.to_string(col, type_=sqltypes.String) 118 | 119 | def bind_processor(self, dialect): 120 | def process(value): 121 | if value is None: 122 | return None 123 | if isinstance(value, set): 124 | return ','.join(str(x) for x in sorted(value)) 125 | return str(value) 126 | return process 127 | 128 | def result_processor(self, dialect, coltype): 129 | def process(value): 130 | if value is None: 131 | return None 132 | return set(int(x) for x in value.split(',') if x) 133 | return process 134 | 135 | 136 | class GEOMETRY(sqltypes.TypeEngine): 137 | __visit_name__ = "GEOMETRY" 138 | 139 | def __init__(self, srid=None): 140 | super(GEOMETRY, self).__init__() 141 | self.srid = srid 142 | 143 | 144 | 145 | class GEOGRAPHY(sqltypes.TypeEngine): 146 | __visit_name__ = "GEOGRAPHY" 147 | native = True 148 | 149 | def __init__(self, srid=None): 150 | super(GEOGRAPHY, self).__init__() 151 | self.srid = srid 152 | 153 | 154 | -------------------------------------------------------------------------------- /tests/test_table_options.py: -------------------------------------------------------------------------------- 1 | 2 | from sqlalchemy.testing import config, fixture, fixtures, util 3 | from sqlalchemy.testing.assertions import AssertsCompiledSQL 4 | from sqlalchemy import Table, Column, Integer, String, func, MetaData, schema, cast 5 | 6 | 7 | class CompileDatabendTableOptionsTest(fixtures.TestBase, AssertsCompiledSQL): 8 | 9 | __only_on__ = "databend" 10 | 11 | def test_create_table_transient_on(self): 12 | m = MetaData() 13 | tbl = Table( 14 | 'atable', m, Column("id", Integer), 15 | databend_transient=True, 16 | ) 17 | self.assert_compile( 18 | schema.CreateTable(tbl), 19 | "CREATE TRANSIENT TABLE atable (id INTEGER)") 20 | 21 | def test_create_table_transient_off(self): 22 | m = MetaData() 23 | tbl = Table( 24 | 'atable', m, Column("id", Integer), 25 | databend_transient=False, 26 | ) 27 | self.assert_compile( 28 | schema.CreateTable(tbl), 29 | "CREATE TABLE atable (id INTEGER)") 30 | 31 | def test_create_table_engine(self): 32 | m = MetaData() 33 | tbl = Table( 34 | 'atable', m, Column("id", Integer), 35 | databend_engine='Memory', 36 | ) 37 | self.assert_compile( 38 | schema.CreateTable(tbl), 39 | "CREATE TABLE atable (id INTEGER) ENGINE=Memory") 40 | 41 | def test_create_table_cluster_by_column_str(self): 42 | m = MetaData() 43 | tbl = Table( 44 | 'atable', m, Column("id", Integer), 45 | databend_cluster_by='id', 46 | ) 47 | self.assert_compile( 48 | schema.CreateTable(tbl), 49 | "CREATE TABLE atable (id INTEGER) CLUSTER BY ( id )") 50 | 51 | def test_create_table_cluster_by_column_strs(self): 52 | m = MetaData() 53 | tbl = Table( 54 | 'atable', m, Column("id", Integer), Column("Name", String), 55 | databend_cluster_by=['id', 'Name'], 56 | ) 57 | self.assert_compile( 58 | schema.CreateTable(tbl), 59 | "CREATE TABLE atable (id INTEGER, \"Name\" VARCHAR) CLUSTER BY ( id, \"Name\" )") 60 | 61 | def test_create_table_cluster_by_column_object(self): 62 | m = MetaData() 63 | c = Column("id", Integer) 64 | tbl = Table( 65 | 'atable', m, c, 66 | databend_cluster_by=[c], 67 | ) 68 | self.assert_compile( 69 | schema.CreateTable(tbl), 70 | "CREATE TABLE atable (id INTEGER) CLUSTER BY ( id )") 71 | 72 | def test_create_table_cluster_by_column_objects(self): 73 | m = MetaData() 74 | c = Column("id", Integer) 75 | c2 = Column("Name", String) 76 | tbl = Table( 77 | 'atable', m, c, c2, 78 | databend_cluster_by=[c, c2], 79 | ) 80 | self.assert_compile( 81 | schema.CreateTable(tbl), 82 | "CREATE TABLE atable (id INTEGER, \"Name\" VARCHAR) CLUSTER BY ( id, \"Name\" )") 83 | 84 | def test_create_table_cluster_by_column_expr(self): 85 | m = MetaData() 86 | c = Column("id", Integer) 87 | c2 = Column("Name", String) 88 | tbl = Table( 89 | 'atable', m, c, c2, 90 | databend_cluster_by=[cast(c, String), c2], 91 | ) 92 | self.assert_compile( 93 | schema.CreateTable(tbl), 94 | "CREATE TABLE atable (id INTEGER, \"Name\" VARCHAR) CLUSTER BY ( CAST(id AS VARCHAR), \"Name\" )") 95 | 96 | def test_create_table_cluster_by_str(self): 97 | m = MetaData() 98 | c = Column("id", Integer) 99 | c2 = Column("Name", String) 100 | tbl = Table( 101 | 'atable', m, c, c2, 102 | databend_cluster_by="CAST(id AS VARCHAR), \"Name\"", 103 | ) 104 | self.assert_compile( 105 | schema.CreateTable(tbl), 106 | "CREATE TABLE atable (id INTEGER, \"Name\" VARCHAR) CLUSTER BY ( CAST(id AS VARCHAR), \"Name\" )") 107 | 108 | #ToDo 109 | # def test_create_table_with_options(self): 110 | # m = MetaData() 111 | # tbl = Table( 112 | # 'atable', m, Column("id", Integer), 113 | # databend_engine_options=( 114 | # ("compression", "snappy"), 115 | # ("storage_format", "parquet"), 116 | # )) 117 | # self.assert_compile( 118 | # schema.CreateTable(tbl), 119 | # "CREATE TABLE atable (id INTEGER)COMPRESSION=\"snappy\" STORAGE_FORMAT=\"parquet\"") 120 | 121 | 122 | class ReflectDatabendTableOptionsTest(fixtures.TablesTest): 123 | __backend__ = True 124 | __only_on__ = "databend" 125 | 126 | # 'once', 'each', None 127 | run_inserts = "None" 128 | 129 | # 'each', None 130 | run_deletes = "None" 131 | 132 | @classmethod 133 | def define_tables(cls, metadata): 134 | Table( 135 | "t2_engine", 136 | metadata, 137 | Column("id", Integer, primary_key=True), 138 | Column("Name", String), 139 | databend_engine="Memory", 140 | ) 141 | c2 = Column("id", Integer, primary_key=True) 142 | Table( 143 | "t2_cluster_by_column", 144 | metadata, 145 | c2, 146 | Column("Name", String), 147 | databend_cluster_by=[c2, "Name"], 148 | ) 149 | c3 = Column("id", Integer, primary_key=True) 150 | Table( 151 | "t3_cluster_by_expr", 152 | metadata, 153 | c3, 154 | Column("Name", String), 155 | databend_cluster_by=[cast(c3, String), "Name"], 156 | ) 157 | c4 = Column("id", Integer, primary_key=True) 158 | Table( 159 | "t4_cluster_by_str", 160 | metadata, 161 | c4, 162 | Column("Name", String), 163 | databend_cluster_by='CAST(id AS STRING), "Name"', 164 | ) 165 | 166 | def test_reflect_table_engine(self): 167 | m2 = MetaData() 168 | t1_ref = Table( 169 | "t2_engine", m2, autoload_with=config.db 170 | ) 171 | assert t1_ref.dialect_options['databend']['engine'] == 'MEMORY' 172 | 173 | def test_reflect_table_cluster_by_column(self): 174 | m2 = MetaData() 175 | t2_ref = Table( 176 | "t2_cluster_by_column", m2, autoload_with=config.db 177 | ) 178 | assert t2_ref.dialect_options['databend']['cluster_by'] == 'id, "Name"' 179 | 180 | def test_reflect_table_cluster_by_expr(self): 181 | m2 = MetaData() 182 | t3_ref = Table( 183 | "t3_cluster_by_expr", m2, autoload_with=config.db 184 | ) 185 | assert t3_ref.dialect_options['databend']['cluster_by'] == 'CAST(id AS STRING), "Name"' 186 | 187 | def test_reflect_table_cluster_by_str(self): 188 | m2 = MetaData() 189 | t4_ref = Table( 190 | "t4_cluster_by_str", m2, autoload_with=config.db 191 | ) 192 | assert t4_ref.dialect_options['databend']['cluster_by'] == 'CAST(id AS STRING), "Name"' -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | databend-sqlalchemy 2 | =================== 3 | 4 | Databend dialect for SQLAlchemy. 5 | 6 | Installation 7 | ------------ 8 | 9 | The package is installable through PIP:: 10 | 11 | pip install databend-sqlalchemy 12 | 13 | Usage 14 | ----- 15 | 16 | The DSN format is similar to that of regular Postgres:: 17 | 18 | from sqlalchemy import create_engine, text 19 | from sqlalchemy.engine.base import Connection, Engine 20 | engine = create_engine( 21 | f"databend://{username}:{password}@{host_port_name}/{database_name}?sslmode=disable" 22 | ) 23 | connection = engine.connect() 24 | result = connection.execute(text("SELECT 1")) 25 | assert len(result.fetchall()) == 1 26 | 27 | import connector 28 | cursor = connector.connect('databend://root:@localhost:8000?sslmode=disable').cursor() 29 | cursor.execute('SELECT * FROM test') 30 | # print(cursor.fetchone()) 31 | # print(cursor.fetchall()) 32 | for row in cursor: 33 | print(row) 34 | 35 | 36 | Merge Command Support 37 | --------------------- 38 | 39 | Databend SQLAlchemy supports upserts via its `Merge` custom expression. 40 | See [Merge](https://docs.databend.com/sql/sql-commands/dml/dml-merge) for full documentation. 41 | 42 | The Merge command can be used as below:: 43 | 44 | from sqlalchemy.orm import sessionmaker 45 | from sqlalchemy import MetaData, create_engine 46 | from databend_sqlalchemy.databend_dialect import Merge 47 | 48 | engine = create_engine(db.url, echo=False) 49 | session = sessionmaker(bind=engine)() 50 | connection = engine.connect() 51 | 52 | meta = MetaData() 53 | meta.reflect(bind=session.bind) 54 | t1 = meta.tables['t1'] 55 | t2 = meta.tables['t2'] 56 | 57 | merge = Merge(target=t1, source=t2, on=t1.c.t1key == t2.c.t2key) 58 | merge.when_matched_then_delete().where(t2.c.marked == 1) 59 | merge.when_matched_then_update().where(t2.c.isnewstatus == 1).values(val = t2.c.newval, status=t2.c.newstatus) 60 | merge.when_matched_then_update().values(val=t2.c.newval) 61 | merge.when_not_matched_then_insert().values(val=t2.c.newval, status=t2.c.newstatus) 62 | connection.execute(merge) 63 | 64 | 65 | Copy Into Command Support 66 | --------------------- 67 | 68 | Databend SQLAlchemy supports copy into operations through it's CopyIntoTable and CopyIntoLocation methods 69 | See [CopyIntoLocation](https://docs.databend.com/sql/sql-commands/dml/dml-copy-into-location) or [CopyIntoTable](https://docs.databend.com/sql/sql-commands/dml/dml-copy-into-table) for full documentation. 70 | 71 | The CopyIntoTable command can be used as below:: 72 | 73 | from sqlalchemy.orm import sessionmaker 74 | from sqlalchemy import MetaData, create_engine 75 | from databend_sqlalchemy import ( 76 | CopyIntoTable, GoogleCloudStorage, ParquetFormat, CopyIntoTableOptions, 77 | FileColumnClause, CSVFormat, 78 | ) 79 | 80 | engine = create_engine(db.url, echo=False) 81 | session = sessionmaker(bind=engine)() 82 | connection = engine.connect() 83 | 84 | meta = MetaData() 85 | meta.reflect(bind=session.bind) 86 | t1 = meta.tables['t1'] 87 | t2 = meta.tables['t2'] 88 | gcs_private_key = 'full_gcs_json_private_key' 89 | case_sensitive_columns = True 90 | 91 | copy_into = CopyIntoTable( 92 | target=t1, 93 | from_=GoogleCloudStorage( 94 | uri='gcs://bucket-name/path/to/file', 95 | credentials=base64.b64encode(gcs_private_key.encode()).decode(), 96 | ), 97 | file_format=ParquetFormat(), 98 | options=CopyIntoTableOptions( 99 | force=True, 100 | column_match_mode='CASE_SENSITIVE' if case_sensitive_columns else None, 101 | ) 102 | ) 103 | result = connection.execute(copy_into) 104 | result.fetchall() # always call fetchall() to ensure the cursor executes to completion 105 | 106 | # More involved example with column selection clause that can be altered to perform operations on the columns during import. 107 | 108 | copy_into = CopyIntoTable( 109 | target=t2, 110 | from_=FileColumnClause( 111 | columns=', '.join([ 112 | f'${index + 1}' 113 | for index, column in enumerate(t2.columns) 114 | ]), 115 | from_=GoogleCloudStorage( 116 | uri='gcs://bucket-name/path/to/file', 117 | credentials=base64.b64encode(gcs_private_key.encode()).decode(), 118 | ) 119 | ), 120 | pattern='*.*', 121 | file_format=CSVFormat( 122 | record_delimiter='\n', 123 | field_delimiter=',', 124 | quote='"', 125 | escape='', 126 | skip_header=1, 127 | empty_field_as='NULL', 128 | compression=Compression.AUTO, 129 | ), 130 | options=CopyIntoTableOptions( 131 | force=True, 132 | ) 133 | ) 134 | result = connection.execute(copy_into) 135 | result.fetchall() # always call fetchall() to ensure the cursor executes to completion 136 | 137 | The CopyIntoLocation command can be used as below:: 138 | 139 | from sqlalchemy.orm import sessionmaker 140 | from sqlalchemy import MetaData, create_engine 141 | from databend_sqlalchemy import ( 142 | CopyIntoLocation, GoogleCloudStorage, ParquetFormat, CopyIntoLocationOptions, 143 | ) 144 | 145 | engine = create_engine(db.url, echo=False) 146 | session = sessionmaker(bind=engine)() 147 | connection = engine.connect() 148 | 149 | meta = MetaData() 150 | meta.reflect(bind=session.bind) 151 | t1 = meta.tables['t1'] 152 | gcs_private_key = 'full_gcs_json_private_key' 153 | 154 | copy_into = CopyIntoLocation( 155 | target=GoogleCloudStorage( 156 | uri='gcs://bucket-name/path/to/target_file', 157 | credentials=base64.b64encode(gcs_private_key.encode()).decode(), 158 | ), 159 | from_=select(t1).where(t1.c['col1'] == 1), 160 | file_format=ParquetFormat(), 161 | options=CopyIntoLocationOptions( 162 | single=True, 163 | overwrite=True, 164 | include_query_id=False, 165 | use_raw_path=True, 166 | ) 167 | ) 168 | result = connection.execute(copy_into) 169 | result.fetchall() # always call fetchall() to ensure the cursor executes to completion 170 | 171 | Table Options 172 | --------------------- 173 | 174 | Databend SQLAlchemy supports databend specific table options for Engine, Cluster Keys and Transient tables 175 | 176 | The table options can be used as below:: 177 | 178 | from sqlalchemy import Table, Column 179 | from sqlalchemy import MetaData, create_engine 180 | 181 | engine = create_engine(db.url, echo=False) 182 | 183 | meta = MetaData() 184 | # Example of Transient Table 185 | t_transient = Table( 186 | "t_transient", 187 | meta, 188 | Column("c1", Integer), 189 | databend_transient=True, 190 | ) 191 | 192 | # Example of Engine 193 | t_engine = Table( 194 | "t_engine", 195 | meta, 196 | Column("c1", Integer), 197 | databend_engine='Memory', 198 | ) 199 | 200 | # Examples of Table with Cluster Keys 201 | t_cluster_1 = Table( 202 | "t_cluster_1", 203 | meta, 204 | Column("c1", Integer), 205 | databend_cluster_by=[c1], 206 | ) 207 | # 208 | c = Column("id", Integer) 209 | c2 = Column("Name", String) 210 | t_cluster_2 = Table( 211 | 't_cluster_2', 212 | meta, 213 | c, 214 | c2, 215 | databend_cluster_by=[cast(c, String), c2], 216 | ) 217 | 218 | meta.create_all(engine) 219 | 220 | 221 | 222 | Compatibility 223 | --------------- 224 | 225 | - If databend version >= v0.9.0 or later, you need to use databend-sqlalchemy version >= v0.1.0. 226 | - The databend-sqlalchemy use [databend-py](https://github.com/databendlabs/databend-py) as internal driver when version < v0.4.0, but when version >= v0.4.0 it use [databend driver python binding](https://github.com/databendlabs/bendsql/blob/main/bindings/python/README.md) as internal driver. The only difference between the two is that the connection parameters provided in the DSN are different. When using the corresponding version, you should refer to the connection parameters provided by the corresponding Driver. 227 | -------------------------------------------------------------------------------- /tests/test_copy_into.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from sqlalchemy.testing import config, fixture, fixtures, eq_ 4 | from sqlalchemy.testing.assertions import AssertsCompiledSQL 5 | from sqlalchemy import ( 6 | Table, 7 | Column, 8 | Integer, 9 | String, 10 | func, 11 | MetaData, 12 | schema, 13 | cast, 14 | literal_column, 15 | text, 16 | ) 17 | 18 | from databend_sqlalchemy import ( 19 | CopyIntoTable, 20 | CopyIntoLocation, 21 | CopyIntoTableOptions, 22 | CopyIntoLocationOptions, 23 | CSVFormat, 24 | ParquetFormat, 25 | GoogleCloudStorage, 26 | Compression, 27 | FileColumnClause, 28 | StageClause, 29 | ) 30 | import sqlalchemy 31 | from packaging import version 32 | 33 | 34 | class CompileDatabendCopyIntoTableTest(fixtures.TestBase, AssertsCompiledSQL): 35 | 36 | __only_on__ = "databend" 37 | 38 | def test_copy_into_table(self): 39 | m = MetaData() 40 | tbl = Table( 41 | "atable", 42 | m, 43 | Column("id", Integer), 44 | schema="test_schema", 45 | ) 46 | 47 | copy_into = CopyIntoTable( 48 | target=tbl, 49 | from_=GoogleCloudStorage( 50 | uri="gcs://some-bucket/a/path/to/files", 51 | credentials="XYZ", 52 | ), 53 | # files='', 54 | # pattern='', 55 | file_format=CSVFormat( 56 | record_delimiter="\n", 57 | field_delimiter=",", 58 | quote='"', 59 | # escape='\\', 60 | # skip_header=1, 61 | # nan_display='' 62 | # null_display='', 63 | error_on_column_mismatch=False, 64 | # empty_field_as='STRING', 65 | output_header=True, 66 | # binary_format='', 67 | compression=Compression.GZIP, 68 | ), 69 | options=CopyIntoTableOptions( 70 | size_limit=None, 71 | purge=None, 72 | force=None, 73 | disable_variant_check=None, 74 | on_error=None, 75 | max_files=None, 76 | return_failed_only=None, 77 | column_match_mode=None, 78 | ), 79 | ) 80 | 81 | self.assert_compile( 82 | copy_into, 83 | ( 84 | "COPY INTO test_schema.atable" 85 | " FROM 'gcs://some-bucket/a/path/to/files' " 86 | "CONNECTION = (" 87 | " ENDPOINT_URL = 'https://storage.googleapis.com' " 88 | " CREDENTIAL = 'XYZ' " 89 | ")" 90 | " FILE_FORMAT = (TYPE = CSV, " 91 | "RECORD_DELIMITER = '\\n', FIELD_DELIMITER = ',', QUOTE = '\"', OUTPUT_HEADER = TRUE, COMPRESSION = GZIP) " 92 | ), 93 | ) 94 | 95 | def test_copy_into_table_sub_select_string_columns(self): 96 | m = MetaData() 97 | tbl = Table( 98 | "atable", 99 | m, 100 | Column("id", Integer), 101 | schema="test_schema", 102 | ) 103 | 104 | copy_into = CopyIntoTable( 105 | target=tbl, 106 | from_=FileColumnClause( 107 | columns="$1, $2, $3", 108 | from_=GoogleCloudStorage( 109 | uri="gcs://some-bucket/a/path/to/files", 110 | credentials="XYZ", 111 | ), 112 | ), 113 | file_format=CSVFormat(), 114 | ) 115 | 116 | self.assert_compile( 117 | copy_into, 118 | ( 119 | "COPY INTO test_schema.atable" 120 | " FROM (SELECT $1, $2, $3" 121 | " FROM 'gcs://some-bucket/a/path/to/files' " 122 | "CONNECTION = (" 123 | " ENDPOINT_URL = 'https://storage.googleapis.com' " 124 | " CREDENTIAL = 'XYZ' " 125 | ")" 126 | ") FILE_FORMAT = (TYPE = CSV)" 127 | ), 128 | ) 129 | 130 | def test_copy_into_table_sub_select_column_clauses(self): 131 | m = MetaData() 132 | tbl = Table( 133 | "atable", 134 | m, 135 | Column("id", Integer), 136 | schema="test_schema", 137 | ) 138 | 139 | copy_into = CopyIntoTable( 140 | target=tbl, 141 | from_=FileColumnClause( 142 | columns=[func.IF(literal_column("$1") == "xyz", "NULL", "NOTNULL")], 143 | # columns='$1, $2, $3', 144 | from_=GoogleCloudStorage( 145 | uri="gcs://some-bucket/a/path/to/files", 146 | credentials="XYZ", 147 | ), 148 | ), 149 | file_format=CSVFormat(), 150 | ) 151 | 152 | self.assert_compile( 153 | copy_into, 154 | ( 155 | "COPY INTO test_schema.atable" 156 | " FROM (SELECT IF($1 = %(1_1)s, %(IF_1)s, %(IF_2)s)" 157 | " FROM 'gcs://some-bucket/a/path/to/files' " 158 | "CONNECTION = (" 159 | " ENDPOINT_URL = 'https://storage.googleapis.com' " 160 | " CREDENTIAL = 'XYZ' " 161 | ")" 162 | ") FILE_FORMAT = (TYPE = CSV)" 163 | ), 164 | checkparams={"1_1": "xyz", "IF_1": "NULL", "IF_2": "NOTNULL"}, 165 | ) 166 | 167 | def test_copy_into_table_files(self): 168 | m = MetaData() 169 | tbl = Table( 170 | "atable", 171 | m, 172 | Column("id", Integer), 173 | schema="test_schema", 174 | ) 175 | 176 | copy_into = CopyIntoTable( 177 | target=tbl, 178 | from_=GoogleCloudStorage( 179 | uri="gcs://some-bucket/a/path/to/files", 180 | credentials="XYZ", 181 | ), 182 | files=['one','two','three'], 183 | file_format=CSVFormat(), 184 | ) 185 | 186 | self.assert_compile( 187 | copy_into, 188 | ( 189 | "COPY INTO test_schema.atable" 190 | " FROM 'gcs://some-bucket/a/path/to/files' " 191 | "CONNECTION = (" 192 | " ENDPOINT_URL = 'https://storage.googleapis.com' " 193 | " CREDENTIAL = 'XYZ' " 194 | ") FILES = ('one', 'two', 'three')" 195 | " FILE_FORMAT = (TYPE = CSV)" 196 | ), 197 | ) 198 | 199 | 200 | class CopyIntoResultTest(fixtures.TablesTest): 201 | run_create_tables = "each" 202 | __backend__ = True 203 | 204 | @classmethod 205 | def define_tables(cls, metadata): 206 | Table( 207 | "random_data", 208 | metadata, 209 | Column("id", Integer), 210 | Column("data", String(50)), 211 | databend_engine='Random', 212 | ) 213 | Table( 214 | "loaded", 215 | metadata, 216 | Column("id", Integer), 217 | Column("data", String(50)), 218 | ) 219 | 220 | if version.parse(sqlalchemy.__version__) >= version.parse('2.0.0'): 221 | def test_copy_into_stage_and_table(self, connection): 222 | # create stage 223 | connection.execute(text('CREATE OR REPLACE STAGE mystage')) 224 | # copy into stage from random table limiting 1000 225 | table = self.tables.random_data 226 | query = table.select().limit(1000) 227 | 228 | copy_into = CopyIntoLocation( 229 | target=StageClause( 230 | name='mystage' 231 | ), 232 | from_=query, 233 | file_format=ParquetFormat(), 234 | options=CopyIntoLocationOptions() 235 | ) 236 | r = connection.execute( 237 | copy_into 238 | ) 239 | eq_(r.rowcount, 1000) 240 | copy_into_results = r.context.copy_into_location_results() 241 | eq_(copy_into_results['rows_unloaded'], 1000) 242 | # eq_(copy_into_results['input_bytes'], 16250) # input bytes will differ, the table is random 243 | # eq_(copy_into_results['output_bytes'], 4701) # output bytes differs 244 | 245 | # now copy into table 246 | 247 | copy_into_table = CopyIntoTable( 248 | target=self.tables.loaded, 249 | from_=StageClause( 250 | name='mystage' 251 | ), 252 | file_format=ParquetFormat(), 253 | options=CopyIntoTableOptions() 254 | ) 255 | r = connection.execute( 256 | copy_into_table 257 | ) 258 | eq_(r.rowcount, 1000) 259 | copy_into_table_results = r.context.copy_into_table_results() 260 | assert len(copy_into_table_results) == 1 261 | result = copy_into_table_results[0] 262 | assert result['file'].endswith('.parquet') 263 | eq_(result['rows_loaded'], 1000) 264 | eq_(result['errors_seen'], 0) 265 | eq_(result['first_error'], None) 266 | eq_(result['first_error_line'], None) 267 | 268 | 269 | -------------------------------------------------------------------------------- /tests/unit/test_databend_dialect.py: -------------------------------------------------------------------------------- 1 | import os 2 | from unittest import mock 3 | 4 | import sqlalchemy 5 | from conftest import MockCursor, MockDBApi 6 | from pytest import mark 7 | from sqlalchemy.engine import url 8 | from sqlalchemy.sql import text 9 | 10 | import databend_sqlalchemy # SQLAlchemy package 11 | from databend_sqlalchemy.databend_dialect import ( 12 | DatabendCompiler, 13 | DatabendDialect, 14 | DatabendIdentifierPreparer, 15 | DatabendTypeCompiler, 16 | INTEGER, BOOLEAN, BINARY, 17 | DatabendDate 18 | ) 19 | from databend_sqlalchemy.databend_dialect import dialect as dialect_definition 20 | 21 | 22 | class TestDatabendDialect: 23 | def test_create_dialect(self, dialect: DatabendDialect): 24 | assert issubclass(dialect_definition, DatabendDialect) 25 | assert isinstance(DatabendDialect.dbapi(), type(databend_sqlalchemy)) 26 | assert dialect.name == "databend" 27 | assert dialect.driver == "databend" 28 | assert issubclass(dialect.preparer, DatabendIdentifierPreparer) 29 | assert issubclass(dialect.statement_compiler, DatabendCompiler) 30 | # SQLAlchemy's DefaultDialect creates an instance of 31 | # type_compiler behind the scenes 32 | assert isinstance(dialect.type_compiler, DatabendTypeCompiler) 33 | assert dialect.context == {} 34 | 35 | def test_create_connect_args(self, dialect: DatabendDialect): 36 | u = url.make_url("databend://user:pass@localhost:8000/testdb") 37 | result_list, result_dict = dialect.create_connect_args(u) 38 | assert result_dict["dsn"] == "databend://user:pass@localhost:8000/testdb" 39 | 40 | u = url.make_url("databend://user:pass@host:443/db") 41 | args, kwargs = dialect.create_connect_args(u) 42 | assert args == [] 43 | assert kwargs["dsn"] == "databend://user:pass@host:443/db" 44 | 45 | u = url.make_url("databend://user:pass@host:443/db?warehouse=test&secure=True") 46 | args, kwargs = dialect.create_connect_args(u) 47 | assert args == [] 48 | assert kwargs["dsn"] == "databend://user:pass@host:443/db?warehouse=test&secure=True" 49 | 50 | def test_do_execute( 51 | self, dialect: DatabendDialect, cursor: mock.Mock(spec=MockCursor) 52 | ): 53 | dialect.do_execute(cursor, "SELECT *", None) 54 | cursor.execute.assert_called_once_with("SELECT *", None) 55 | cursor.execute.reset_mock() 56 | dialect.do_execute(cursor, "SELECT *", (1, 22), None) 57 | 58 | def test_table_names( 59 | self, dialect: DatabendDialect, connection: mock.Mock(spec=MockDBApi) 60 | ): 61 | connection.execute.return_value = [ 62 | ("table1",), 63 | ("table2",), 64 | ] 65 | 66 | result = dialect.get_table_names(connection) 67 | assert result == ["table1", "table2"] 68 | connection.execute.assert_called_once() 69 | assert str(connection.execute.call_args[0][0].compile()) == str( 70 | text(""" 71 | select table_name 72 | from information_schema.tables 73 | where table_schema = :schema_name 74 | and engine NOT LIKE '%VIEW%' 75 | """).compile() 76 | ) 77 | assert connection.execute.call_args[0][1] == {'schema_name': None} 78 | connection.execute.reset_mock() 79 | # Test default schema 80 | dialect.default_schema_name = 'some-schema' 81 | result = dialect.get_table_names(connection) 82 | assert result == ["table1", "table2"] 83 | connection.execute.assert_called_once() 84 | assert str(connection.execute.call_args[0][0].compile()) == str( 85 | text(""" 86 | select table_name 87 | from information_schema.tables 88 | where table_schema = :schema_name 89 | and engine NOT LIKE '%VIEW%' 90 | """).compile() 91 | ) 92 | assert connection.execute.call_args[0][1] == {'schema_name': 'some-schema'} 93 | connection.execute.reset_mock() 94 | # Test specified schema 95 | result = dialect.get_table_names(connection, schema="schema") 96 | assert result == ["table1", "table2"] 97 | connection.execute.assert_called_once() 98 | assert str(connection.execute.call_args[0][0].compile()) == str( 99 | text(""" 100 | select table_name 101 | from information_schema.tables 102 | where table_schema = :schema_name 103 | and engine NOT LIKE '%VIEW%' 104 | """).compile() 105 | ) 106 | assert connection.execute.call_args[0][1] == {'schema_name': 'schema'} 107 | 108 | def test_view_names( 109 | self, dialect: DatabendDialect, connection: mock.Mock(spec=MockDBApi) 110 | ): 111 | connection.execute.return_value = [] 112 | assert dialect.get_view_names(connection) == [] 113 | 114 | def test_indexes( 115 | self, dialect: DatabendDialect, connection: mock.Mock(spec=MockDBApi) 116 | ): 117 | assert dialect.get_indexes(connection, "table") == [] 118 | 119 | def test_columns( 120 | self, dialect: DatabendDialect, connection: mock.Mock(spec=MockDBApi) 121 | ): 122 | def multi_column_row(columns): 123 | def getitem(self, idx): 124 | for i, result in enumerate(columns): 125 | if idx == i: 126 | return result 127 | 128 | return mock.Mock(__getitem__=getitem) 129 | 130 | connection.execute.return_value = [ 131 | multi_column_row(["name1", "INT", "YES"]), 132 | multi_column_row(["name2", "date", "NO"]), 133 | multi_column_row(["name3", "boolean", "YES"]), 134 | multi_column_row(["name4", "binary", "YES"]) 135 | ] 136 | 137 | expected_query = """ 138 | select column_name, column_type, is_nullable 139 | from information_schema.columns 140 | where table_name = :table_name 141 | and table_schema = :schema_name 142 | """ 143 | 144 | for call, expected_params in ( 145 | ( 146 | lambda: dialect.get_columns(connection, "table"), 147 | {'table_name': 'table', 'schema_name': None}, 148 | ), 149 | ( 150 | lambda: dialect.get_columns(connection, "table", "schema"), 151 | {'table_name': 'table', 'schema_name': 'schema'}, 152 | ), 153 | ): 154 | result = call() 155 | assert result == [ 156 | { 157 | "name": "name1", 158 | "type": INTEGER, 159 | "nullable": True, 160 | "default": None, 161 | }, 162 | { 163 | "name": "name2", 164 | "type": DatabendDate, 165 | "nullable": False, 166 | "default": None, 167 | }, 168 | { 169 | "name": "name3", 170 | "type": BOOLEAN, 171 | "nullable": True, 172 | "default": None, 173 | }, 174 | { 175 | "name": "name4", 176 | "type": BINARY, 177 | "nullable": True, 178 | "default": None, 179 | }, 180 | ] 181 | connection.execute.assert_called_once() 182 | assert str(connection.execute.call_args[0][0].compile()) == str( 183 | text(expected_query).compile() 184 | ) 185 | assert connection.execute.call_args[0][1] == expected_params 186 | connection.execute.reset_mock() 187 | 188 | 189 | def test_get_is_nullable(): 190 | assert databend_sqlalchemy.databend_dialect.get_is_nullable("YES") 191 | assert not databend_sqlalchemy.databend_dialect.get_is_nullable("NO") 192 | 193 | 194 | def test_types(): 195 | assert databend_sqlalchemy.databend_dialect.CHAR is sqlalchemy.sql.sqltypes.CHAR 196 | assert issubclass(databend_sqlalchemy.databend_dialect.DatabendDate, sqlalchemy.sql.sqltypes.DATE) 197 | assert issubclass( 198 | databend_sqlalchemy.databend_dialect.DatabendDateTime, 199 | sqlalchemy.sql.sqltypes.DATETIME, 200 | ) 201 | assert ( 202 | databend_sqlalchemy.databend_dialect.INTEGER is sqlalchemy.sql.sqltypes.INTEGER 203 | ) 204 | assert databend_sqlalchemy.databend_dialect.BIGINT is sqlalchemy.sql.sqltypes.BIGINT 205 | assert ( 206 | databend_sqlalchemy.databend_dialect.TIMESTAMP 207 | is sqlalchemy.sql.sqltypes.TIMESTAMP 208 | ) 209 | assert ( 210 | databend_sqlalchemy.databend_dialect.VARCHAR is sqlalchemy.sql.sqltypes.VARCHAR 211 | ) 212 | assert ( 213 | databend_sqlalchemy.databend_dialect.BOOLEAN is sqlalchemy.sql.sqltypes.BOOLEAN 214 | ) 215 | assert databend_sqlalchemy.databend_dialect.FLOAT is sqlalchemy.sql.sqltypes.FLOAT 216 | assert issubclass( 217 | databend_sqlalchemy.databend_dialect.ARRAY, sqlalchemy.types.TypeEngine 218 | ) 219 | 220 | 221 | def test_extract_nullable_string(): 222 | types = ["INT", "FLOAT", "Nullable(INT)", "Nullable(Decimal(2,4))", "Nullable(Array(INT))", 223 | "Nullable(Map(String, String))", "Decimal(1,2)"] 224 | expected_types = ["int", "float", "int", "decimal", "array", "map", "decimal"] 225 | i = 0 226 | for t in types: 227 | true_type = databend_sqlalchemy.databend_dialect.extract_nullable_string(t).lower() 228 | assert expected_types[i] == true_type 229 | i += 1 230 | print(true_type) 231 | -------------------------------------------------------------------------------- /databend_sqlalchemy/requirements.py: -------------------------------------------------------------------------------- 1 | 2 | from sqlalchemy.testing.requirements import SuiteRequirements 3 | 4 | from sqlalchemy.testing import exclusions 5 | 6 | 7 | class Requirements(SuiteRequirements): 8 | 9 | @property 10 | def foreign_keys(self): 11 | """Target database must support foreign keys.""" 12 | 13 | return exclusions.closed() # Currently no foreign keys in Databend 14 | 15 | @property 16 | def binary_comparisons(self): 17 | """target database/driver can allow BLOB/BINARY fields to be compared 18 | against a bound parameter value. 19 | """ 20 | return exclusions.closed() # Currently no binary type in Databend 21 | 22 | @property 23 | def binary_literals(self): 24 | """target backend supports simple binary literals, e.g. an 25 | expression like:: 26 | 27 | SELECT CAST('foo' AS BINARY) 28 | 29 | Where ``BINARY`` is the type emitted from :class:`.LargeBinary`, 30 | e.g. it could be ``BLOB`` or similar. 31 | 32 | Basically fails on Oracle. 33 | 34 | """ 35 | return exclusions.closed() # Currently no binary type in Databend 36 | 37 | @property 38 | def comment_reflection(self): 39 | """Indicates if the database support table comment reflection""" 40 | return exclusions.open() 41 | 42 | @property 43 | def comment_reflection_full_unicode(self): 44 | """Indicates if the database support table comment reflection in the 45 | full unicode range, including emoji etc. 46 | """ 47 | return exclusions.open() 48 | 49 | @property 50 | def temporary_tables(self): 51 | """target database supports temporary tables""" 52 | return exclusions.open() 53 | 54 | @property 55 | def temp_table_reflection(self): 56 | return exclusions.closed() 57 | 58 | @property 59 | def self_referential_foreign_keys(self): 60 | """Target database must support self-referential foreign keys.""" 61 | 62 | return exclusions.closed() # Databend does not currently support foreign keys 63 | 64 | @property 65 | def foreign_key_ddl(self): 66 | """Target database must support the DDL phrases for FOREIGN KEY.""" 67 | 68 | return exclusions.closed() # Databend does not currently support foreign keys 69 | 70 | @property 71 | def index_reflection(self): 72 | return exclusions.closed() # Databend does not currently support indexes 73 | 74 | @property 75 | def primary_key_constraint_reflection(self): 76 | return exclusions.closed() # Databend does not currently support primary keys 77 | 78 | @property 79 | def foreign_key_constraint_reflection(self): 80 | return exclusions.closed() # Databend does not currently support foreign keys 81 | 82 | @property 83 | def unique_constraint_reflection(self): 84 | """target dialect supports reflection of unique constraints""" 85 | return exclusions.closed() # Databend does not currently support unique constraints 86 | 87 | @property 88 | def duplicate_key_raises_integrity_error(self): 89 | """target dialect raises IntegrityError when reporting an INSERT 90 | with a primary key violation. (hint: it should) 91 | 92 | """ 93 | return exclusions.closed() # Databend does not currently support primary keys 94 | 95 | @property 96 | def sql_expression_limit_offset(self): 97 | """target database can render LIMIT and/or OFFSET with a complete 98 | SQL expression, such as one that uses the addition operator. 99 | parameter 100 | """ 101 | 102 | return exclusions.closed() # Databend does not currently support expressions in limit/offset 103 | 104 | @property 105 | def autoincrement_without_sequence(self): 106 | """If autoincrement=True on a column does not require an explicit 107 | sequence. This should be false only for oracle. 108 | """ 109 | return exclusions.closed() 110 | 111 | @property 112 | def datetime_timezone(self): 113 | """target dialect supports representation of Python 114 | datetime.datetime() with tzinfo with DateTime(timezone=True).""" 115 | 116 | return exclusions.closed() 117 | 118 | # @property 119 | # def datetime_implicit_bound(self): 120 | # """target dialect when given a datetime object will bind it such 121 | # that the database server knows the object is a datetime, and not 122 | # a plain string. 123 | # 124 | # """ 125 | # return exclusions.closed() # `SELECT '2012-10-15 12:57:18' AS thing` does not yield a timestamp in Databend 126 | 127 | @property 128 | def datetime_microseconds(self): 129 | """target dialect supports representation of Python 130 | datetime.datetime() with microsecond objects.""" 131 | 132 | return exclusions.open() 133 | 134 | @property 135 | def timestamp_microseconds(self): 136 | """target dialect supports representation of Python 137 | datetime.datetime() with microsecond objects but only 138 | if TIMESTAMP is used.""" 139 | return exclusions.open() 140 | 141 | @property 142 | def time(self): 143 | """target dialect supports representation of Python 144 | datetime.time() objects.""" 145 | 146 | return exclusions.open() 147 | 148 | @property 149 | def time_microseconds(self): 150 | """target dialect supports representation of Python 151 | datetime.time() with microsecond objects.""" 152 | 153 | return exclusions.open() 154 | 155 | @property 156 | def time_timezone(self): 157 | """target dialect supports representation of Python 158 | datetime.time() with tzinfo with Time(timezone=True).""" 159 | 160 | return exclusions.closed() 161 | 162 | @property 163 | def datetime_interval(self): 164 | """target dialect supports representation of Python 165 | datetime.timedelta().""" 166 | 167 | return exclusions.open() 168 | 169 | @property 170 | def autoincrement_insert(self): 171 | """target platform generates new surrogate integer primary key values 172 | when insert() is executed, excluding the pk column.""" 173 | 174 | return exclusions.closed() 175 | 176 | @property 177 | def views(self): 178 | """Target database must support VIEWs.""" 179 | 180 | return exclusions.open() 181 | 182 | @property 183 | def unicode_data(self): 184 | """Target database/dialect must support Python unicode objects with 185 | non-ASCII characters represented, delivered as bound parameters 186 | as well as in result rows. 187 | 188 | """ 189 | return exclusions.open() 190 | 191 | @property 192 | def unicode_ddl(self): 193 | """Target driver must support some degree of non-ascii symbol 194 | names. 195 | """ 196 | return exclusions.open() 197 | 198 | @property 199 | def precision_generic_float_type(self): 200 | """target backend will return native floating point numbers with at 201 | least seven decimal places when using the generic Float type. 202 | 203 | """ 204 | return exclusions.closed() #ToDo - I couldn't get the test for this one working, not sure where the issue is - AssertionError: {Decimal('15.7563829')} != {Decimal('15.7563827')} 205 | 206 | @property 207 | def precision_numerics_many_significant_digits(self): 208 | """target backend supports values with many digits on both sides, 209 | such as 319438950232418390.273596, 87673.594069654243 210 | 211 | """ 212 | return exclusions.closed() 213 | 214 | @property 215 | def array_type(self): 216 | return exclusions.closed() 217 | 218 | @property 219 | def float_is_numeric(self): 220 | """target backend uses Numeric for Float/Dual""" 221 | 222 | return exclusions.closed() 223 | 224 | @property 225 | def json_type(self): 226 | """target platform implements a native JSON type.""" 227 | 228 | return exclusions.closed() # ToDo - not quite ready to turn on yet, null values are not handled correctly https://github.com/databendlabs/databend/issues/17433 229 | 230 | @property 231 | def reflect_table_options(self): 232 | """Target database must support reflecting table_options.""" 233 | return exclusions.open() 234 | 235 | @property 236 | def ctes(self): 237 | """Target database supports CTEs""" 238 | return exclusions.open() 239 | 240 | @property 241 | def ctes_with_update_delete(self): 242 | """target database supports CTES that ride on top of a normal UPDATE 243 | or DELETE statement which refers to the CTE in a correlated subquery. 244 | 245 | """ 246 | return exclusions.open() 247 | 248 | @property 249 | def update_from(self): 250 | """Target must support UPDATE..FROM syntax""" 251 | return exclusions.closed() 252 | 253 | 254 | @property 255 | def delete_from(self): 256 | """Target must support DELETE FROM..FROM or DELETE..USING syntax""" 257 | return exclusions.closed() 258 | 259 | @property 260 | def table_value_constructor(self): 261 | """Database / dialect supports a query like: 262 | 263 | .. sourcecode:: sql 264 | 265 | SELECT * FROM VALUES ( (c1, c2), (c1, c2), ...) 266 | AS some_table(col1, col2) 267 | 268 | SQLAlchemy generates this with the :func:`_sql.values` function. 269 | 270 | """ 271 | return exclusions.open() 272 | 273 | @property 274 | def window_functions(self): 275 | """Target database must support window functions.""" 276 | return exclusions.open() 277 | -------------------------------------------------------------------------------- /databend_sqlalchemy/connector.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # 3 | # See http://www.python.org/dev/peps/pep-0249/ 4 | # 5 | # Many docstrings in this file are based on the PEP, which is in the public domain. 6 | import decimal 7 | import re 8 | import json 9 | from datetime import datetime, date, time, timedelta 10 | from databend_sqlalchemy.errors import Error, NotSupportedError 11 | 12 | from databend_driver import BlockingDatabendClient 13 | 14 | # PEP 249 module globals 15 | apilevel = "2.0" 16 | threadsafety = 2 # Threads may share the module and connections. 17 | paramstyle = "pyformat" # Python extended format codes, e.g. ...WHERE name=%(name)s 18 | Binary = bytes # to satisfy dialect.dbapi.Binary 19 | 20 | 21 | class ParamEscaper: 22 | def escape_args(self, parameters): 23 | if isinstance(parameters, dict): 24 | return {k: self.escape_item(v) for k, v in parameters.items()} 25 | elif isinstance(parameters, (list, tuple)): 26 | return tuple(self.escape_item(x) for x in parameters) 27 | else: 28 | raise Exception("Unsupported param format: {}".format(parameters)) 29 | 30 | def escape_number(self, item): 31 | return item 32 | 33 | def escape_string(self, item): 34 | # Need to decode UTF-8 because of old sqlalchemy. 35 | # Newer SQLAlchemy checks dialect.supports_unicode_binds before encoding Unicode strings 36 | # as byte strings. The old version always encodes Unicode as byte strings, which breaks 37 | # string formatting here. 38 | if isinstance(item, bytes): 39 | item = item.decode("utf-8") 40 | return "'{}'".format( 41 | item.replace("\\", "\\\\").replace("'", "\\'").replace("%", "%%") 42 | ) 43 | 44 | def escape_item(self, item): 45 | if item is None: 46 | return "NULL" 47 | elif isinstance(item, (int, float)): 48 | return self.escape_number(item) 49 | elif isinstance(item, decimal.Decimal): 50 | return self.escape_number(item) 51 | elif isinstance(item, timedelta): 52 | return self.escape_string(f"{item.total_seconds()} seconds") + "::interval" 53 | elif isinstance(item, time): 54 | # N.B. Date here must match date in DatabendTime.literal_processor - 1970-01-01 55 | return self.escape_string(item.strftime("1970-01-01 %H:%M:%S.%f")) + "::timestamp" 56 | elif isinstance(item, datetime): 57 | return self.escape_string(item.strftime("%Y-%m-%d %H:%M:%S.%f")) + "::timestamp" 58 | elif isinstance(item, date): 59 | return self.escape_string(item.strftime("%Y-%m-%d")) + "::date" 60 | elif isinstance(item, dict): 61 | return self.escape_string(f'parse_json({json.dumps(item)})') 62 | else: 63 | return self.escape_string(item) 64 | 65 | 66 | _escaper = ParamEscaper() 67 | 68 | 69 | # Patch ORM library 70 | @classmethod 71 | def create_ad_hoc_field(cls, db_type): 72 | # Enums 73 | if db_type.startswith("Enum"): 74 | db_type = "String" # enum.Eum is not comparable 75 | # Arrays 76 | if db_type.startswith("Array"): 77 | return "Array" 78 | # FixedString 79 | if db_type.startswith("FixedString"): 80 | db_type = "String" 81 | 82 | if db_type == "LowCardinality(String)": 83 | db_type = "String" 84 | 85 | if db_type.startswith("DateTime"): 86 | db_type = "DateTime" 87 | 88 | if db_type.startswith("Nullable"): 89 | return "Nullable" 90 | 91 | 92 | # 93 | # Connector interface 94 | # 95 | 96 | 97 | def connect(*args, **kwargs): 98 | return Connection(*args, **kwargs) 99 | 100 | 101 | class Connection: 102 | """ 103 | These objects are small stateless factories for cursors, which do all the real work. 104 | """ 105 | 106 | def __init__(self, dsn="databend://root:@localhost:8000/?sslmode=disable"): 107 | self.client = BlockingDatabendClient(dsn) 108 | 109 | def close(self): 110 | pass 111 | 112 | def commit(self): 113 | pass 114 | 115 | def cursor(self): 116 | return Cursor(self.client.cursor()) 117 | 118 | def rollback(self): 119 | raise NotSupportedError("Transactions are not supported") # pragma: no cover 120 | 121 | 122 | class Cursor: 123 | """These objects represent a database cursor, which is used to manage the context of a fetch 124 | operation. 125 | 126 | Cursors are not isolated, i.e., any changes done to the database by a cursor are immediately 127 | visible by other cursors or connections. 128 | """ 129 | 130 | def __init__(self, conn): 131 | self.inner = conn 132 | 133 | @property 134 | def rowcount(self): 135 | """By default, return -1 to indicate that this is not supported.""" 136 | return -1 137 | 138 | @property 139 | def description(self): 140 | """This read-only attribute is a sequence of 7-item sequences. 141 | 142 | Each of these sequences contains information describing one result column: 143 | 144 | - name 145 | - type_code 146 | - display_size (None in current implementation) 147 | - internal_size (None in current implementation) 148 | - precision (None in current implementation) 149 | - scale (None in current implementation) 150 | - null_ok (always True in current implementation) 151 | 152 | The ``type_code`` can be interpreted by comparing it to the Type Objects specified in the 153 | section below. 154 | """ 155 | try: 156 | return self.inner.description 157 | except Exception as e: 158 | raise Error(str(e)) from e 159 | 160 | def close(self): 161 | try: 162 | self.inner.close() 163 | except Exception as e: 164 | raise Error(str(e)) from e 165 | 166 | def mogrify(self, query, parameters): 167 | if parameters: 168 | query = query % _escaper.escape_args(parameters) 169 | return query 170 | 171 | def execute(self, operation, parameters=None): 172 | """Prepare and execute a database operation (query or command).""" 173 | 174 | # ToDo - Fix this, which is preventing the execution of blank DDL such as CREATE INDEX statements which aren't currently supported 175 | # Seems hard to fix when statements are coming from metadata.create_all() 176 | if not operation: 177 | return 178 | 179 | try: 180 | query = self.mogrify(operation, parameters) 181 | query = query.replace("%%", "%") 182 | return self.inner.execute(query) 183 | except Exception as e: 184 | raise Error(str(e)) from e 185 | 186 | def executemany(self, operation, seq_of_parameters): 187 | """Prepare a database operation (query or command) and then execute it against all parameter 188 | sequences or mappings found in the sequence ``seq_of_parameters``. 189 | 190 | Only the final result set is retained. 191 | 192 | Return values are not defined. 193 | """ 194 | values_list = [] 195 | RE_INSERT_VALUES = re.compile( 196 | r"\s*((?:INSERT|REPLACE)\s.+\sVALUES?\s*)" 197 | + r"(\(\s*(?:%s|%\(.+\)s)\s*(?:,\s*(?:%s|%\(.+\)s)\s*)*\))" 198 | + r"(\s*(?:ON DUPLICATE.*)?);?\s*\Z", 199 | re.IGNORECASE | re.DOTALL, 200 | ) 201 | 202 | m = RE_INSERT_VALUES.match(operation) 203 | if m: 204 | try: 205 | q_prefix = m.group(1) 206 | q_values = m.group(2).rstrip() 207 | 208 | for parameters in seq_of_parameters: 209 | values_list.append(q_values % _escaper.escape_args(parameters)) 210 | query = "{} {};".format(q_prefix, ",".join(values_list)) 211 | return self.inner.execute(query) 212 | except Exception as e: 213 | # We have to raise dbAPI error 214 | raise Error(str(e)) from e 215 | else: 216 | for parameters in seq_of_parameters: 217 | self.execute(operation, parameters) 218 | 219 | def fetchone(self): 220 | """Fetch the next row of a query result set, returning a single sequence, or ``None`` when 221 | no more data is available.""" 222 | try: 223 | row = self.inner.fetchone() 224 | if row is None: 225 | return None 226 | return row.values() 227 | except Exception as e: 228 | raise Error(str(e)) from e 229 | 230 | def fetchmany(self, size=None): 231 | """Fetch the next set of rows of a query result, returning a sequence of sequences (e.g. a 232 | list of tuples). An empty sequence is returned when no more rows are available. 233 | 234 | The number of rows to fetch per call is specified by the parameter. If it is not given, the 235 | cursor's arraysize determines the number of rows to be fetched. The method should try to 236 | fetch as many rows as indicated by the size parameter. If this is not possible due to the 237 | specified number of rows not being available, fewer rows may be returned. 238 | """ 239 | try: 240 | rows = self.inner.fetchmany(size) 241 | return [row.values() for row in rows] 242 | except Exception as e: 243 | raise Error(str(e)) from e 244 | 245 | def fetchall(self): 246 | """Fetch all (remaining) rows of a query result, returning them as a sequence of sequences 247 | (e.g. a list of tuples). 248 | """ 249 | try: 250 | rows = self.inner.fetchall() 251 | return [row.values() for row in rows] 252 | except Exception as e: 253 | raise Error(str(e)) from e 254 | 255 | def __next__(self): 256 | """Return the next row from the currently executing SQL statement using the same semantics 257 | as :py:meth:`fetchone`. A ``StopIteration`` exception is raised when the result set is 258 | exhausted. 259 | """ 260 | try: 261 | return self.inner.__next__() 262 | except StopIteration as e: 263 | raise e 264 | except Exception as e: 265 | raise Error(str(e)) from e 266 | 267 | next = __next__ 268 | 269 | def __iter__(self): 270 | """Return self to make cursors compatible to the iteration protocol.""" 271 | return self 272 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /tests/test_merge.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from sqlalchemy import Column 4 | from sqlalchemy import exc 5 | from sqlalchemy import Integer 6 | from sqlalchemy import schema 7 | from sqlalchemy import sql 8 | from sqlalchemy import String 9 | from sqlalchemy import Table 10 | from sqlalchemy import testing 11 | from sqlalchemy import types as sqltypes 12 | 13 | # from sqlalchemy.dialects.postgresql import insert 14 | from sqlalchemy.testing import config, AssertsCompiledSQL 15 | from sqlalchemy.testing import fixtures 16 | from sqlalchemy.testing.assertions import assert_raises 17 | from sqlalchemy.testing.assertions import eq_ 18 | 19 | from databend_sqlalchemy import Merge 20 | 21 | 22 | class MergeIntoTest(fixtures.TablesTest, AssertsCompiledSQL): 23 | __backend__ = True 24 | run_define_tables = "each" 25 | 26 | @classmethod 27 | def define_tables(cls, metadata): 28 | Table( 29 | "users", 30 | metadata, 31 | Column("id", Integer, primary_key=True), 32 | Column("name", String(50)), 33 | Column("login_email", String(50)), 34 | ) 35 | 36 | Table( 37 | "users_schema", 38 | metadata, 39 | Column("id", Integer, primary_key=True), 40 | Column("name", String(50)), 41 | schema=config.test_schema, 42 | ) 43 | 44 | Table( 45 | "users_xtra", 46 | metadata, 47 | Column("id", Integer, primary_key=True), 48 | Column("name", String(50)), 49 | Column("login_email", String(50)), 50 | ) 51 | 52 | def test_no_action_raises(self, connection): 53 | users = self.tables.users 54 | users_xtra = self.tables.users_xtra 55 | 56 | connection.execute( 57 | users.insert(), dict(id=1, name="name1", login_email="email1") 58 | ) 59 | connection.execute( 60 | users.insert(), dict(id=2, name="name2", login_email="email2") 61 | ) 62 | connection.execute( 63 | users_xtra.insert(), dict(id=1, name="newname1", login_email="newemail1") 64 | ) 65 | connection.execute( 66 | users_xtra.insert(), dict(id=2, name="newname2", login_email="newemail2") 67 | ) 68 | 69 | merge = Merge(users, users_xtra, users.c.id == users_xtra.c.id) 70 | 71 | assert_raises( 72 | exc.DBAPIError, 73 | connection.execute, 74 | merge, 75 | ) 76 | 77 | def test_select_as_source(self, connection): 78 | users = self.tables.users 79 | users_xtra = self.tables.users_xtra 80 | 81 | connection.execute( 82 | users.insert(), dict(id=1, name="name1", login_email="email1") 83 | ) 84 | connection.execute( 85 | users_xtra.insert(), dict(id=1, name="newname1", login_email="newemail1") 86 | ) 87 | connection.execute( 88 | users_xtra.insert(), dict(id=2, name="newname2", login_email="newemail2") 89 | ) 90 | 91 | select_source = users_xtra.select().where(users_xtra.c.id != 99) 92 | merge = Merge( 93 | users, select_source, users.c.id == select_source.selected_columns.id 94 | ) 95 | merge.when_matched_then_update().values( 96 | name=select_source.selected_columns.name 97 | ) 98 | merge.when_not_matched_then_insert() 99 | 100 | result = connection.execute(merge) 101 | eq_( 102 | connection.execute(users.select().order_by(users.c.id)).fetchall(), 103 | [(1, "newname1", "email1"), (2, "newname2", "newemail2")], 104 | ) 105 | 106 | def test_alias_as_source(self, connection): 107 | users = self.tables.users 108 | users_xtra = self.tables.users_xtra 109 | 110 | connection.execute( 111 | users.insert(), dict(id=1, name="name1", login_email="email1") 112 | ) 113 | connection.execute( 114 | users_xtra.insert(), dict(id=1, name="newname1", login_email="newemail1") 115 | ) 116 | connection.execute( 117 | users_xtra.insert(), dict(id=2, name="newname2", login_email="newemail2") 118 | ) 119 | 120 | alias_source = users_xtra.select().where(users_xtra.c.id != 99).alias("x") 121 | merge = Merge(users, alias_source, users.c.id == alias_source.c.id) 122 | merge.when_matched_then_update().values(name=alias_source.c.name) 123 | merge.when_not_matched_then_insert() 124 | 125 | result = connection.execute(merge) 126 | eq_( 127 | connection.execute(users.select().order_by(users.c.id)).fetchall(), 128 | [(1, "newname1", "email1"), (2, "newname2", "newemail2")], 129 | ) 130 | 131 | def test_subquery_as_source(self, connection): 132 | users = self.tables.users 133 | users_xtra = self.tables.users_xtra 134 | 135 | connection.execute( 136 | users.insert(), dict(id=1, name="name1", login_email="email1") 137 | ) 138 | connection.execute( 139 | users_xtra.insert(), dict(id=1, name="newname1", login_email="newemail1") 140 | ) 141 | connection.execute( 142 | users_xtra.insert(), dict(id=2, name="newname2", login_email="newemail2") 143 | ) 144 | 145 | subquery_source = users_xtra.select().where(users_xtra.c.id != 99).subquery() 146 | merge = Merge(users, subquery_source, users.c.id == subquery_source.c.id) 147 | merge.when_matched_then_update().values(name=subquery_source.c.name) 148 | merge.when_not_matched_then_insert() 149 | 150 | result = connection.execute(merge) 151 | eq_( 152 | connection.execute(users.select().order_by(users.c.id)).fetchall(), 153 | [(1, "newname1", "email1"), (2, "newname2", "newemail2")], 154 | ) 155 | 156 | def test_when_not_matched_insert(self, connection): 157 | users = self.tables.users 158 | users_xtra = self.tables.users_xtra 159 | 160 | connection.execute( 161 | users.insert(), dict(id=1, name="name1", login_email="email1") 162 | ) 163 | connection.execute( 164 | users_xtra.insert(), dict(id=1, name="newname1", login_email="newemail1") 165 | ) 166 | connection.execute( 167 | users_xtra.insert(), dict(id=2, name="newname2", login_email="newemail2") 168 | ) 169 | 170 | merge = Merge(users, users_xtra, users.c.id == users_xtra.c.id) 171 | merge.when_not_matched_then_insert() 172 | 173 | self.assert_compile( 174 | merge, 175 | 'MERGE INTO "users" USING (SELECT users_xtra.id AS id, users_xtra.name AS name, users_xtra.login_email AS login_email FROM users_xtra) AS users_xtra ON "users".id = users_xtra.id WHEN NOT MATCHED THEN INSERT *', 176 | ) 177 | 178 | result = connection.execute(merge) 179 | eq_( 180 | connection.execute(users.select().order_by(users.c.id)).fetchall(), 181 | [(1, "name1", "email1"), (2, "newname2", "newemail2")], 182 | ) 183 | 184 | def test_when_matched_update(self, connection): 185 | users = self.tables.users 186 | users_xtra = self.tables.users_xtra 187 | 188 | connection.execute( 189 | users.insert(), dict(id=1, name="name1", login_email="email1") 190 | ) 191 | connection.execute( 192 | users.insert(), dict(id=2, name="name2", login_email="email2") 193 | ) 194 | connection.execute( 195 | users_xtra.insert(), dict(id=1, name="newname1", login_email="newemail1") 196 | ) 197 | 198 | merge = Merge(users, users_xtra, users.c.id == users_xtra.c.id) 199 | merge.when_matched_then_update() 200 | 201 | result = connection.execute(merge) 202 | eq_( 203 | connection.execute(users.select().order_by(users.c.id)).fetchall(), 204 | [(1, "newname1", "newemail1"), (2, "name2", "email2")], 205 | ) 206 | 207 | def test_when_matched_update_column(self, connection): 208 | users = self.tables.users 209 | users_xtra = self.tables.users_xtra 210 | 211 | connection.execute( 212 | users.insert(), dict(id=1, name="name1", login_email="email1") 213 | ) 214 | connection.execute( 215 | users.insert(), dict(id=2, name="name2", login_email="email2") 216 | ) 217 | connection.execute( 218 | users_xtra.insert(), dict(id=1, name="newname1", login_email="newemail1") 219 | ) 220 | 221 | merge = Merge(users, users_xtra, users.c.id == users_xtra.c.id) 222 | merge.when_matched_then_update().values(name=users_xtra.c.name) 223 | 224 | result = connection.execute(merge) 225 | eq_( 226 | connection.execute(users.select().order_by(users.c.id)).fetchall(), 227 | [(1, "newname1", "email1"), (2, "name2", "email2")], 228 | ) 229 | 230 | def test_when_matched_update_criteria(self, connection): 231 | users = self.tables.users 232 | users_xtra = self.tables.users_xtra 233 | 234 | connection.execute( 235 | users.insert(), dict(id=1, name="name1", login_email="email1") 236 | ) 237 | connection.execute( 238 | users.insert(), dict(id=2, name="name2", login_email="email2") 239 | ) 240 | connection.execute( 241 | users_xtra.insert(), dict(id=1, name="newname1", login_email="newemail1") 242 | ) 243 | connection.execute( 244 | users_xtra.insert(), dict(id=2, name="newname2", login_email="newemail2") 245 | ) 246 | 247 | merge = Merge(users, users_xtra, users.c.id == users_xtra.c.id) 248 | merge.when_matched_then_update().where(users_xtra.c.id != 1) 249 | 250 | result = connection.execute(merge) 251 | eq_( 252 | connection.execute(users.select().order_by(users.c.id)).fetchall(), 253 | [(1, "name1", "email1"), (2, "newname2", "newemail2")], 254 | ) 255 | 256 | def test_when_matched_update_criteria_column(self, connection): 257 | users = self.tables.users 258 | users_xtra = self.tables.users_xtra 259 | 260 | connection.execute( 261 | users.insert(), dict(id=1, name="name1", login_email="email1") 262 | ) 263 | connection.execute( 264 | users.insert(), dict(id=2, name="name2", login_email="email2") 265 | ) 266 | connection.execute( 267 | users_xtra.insert(), dict(id=1, name="newname1", login_email="newemail1") 268 | ) 269 | connection.execute( 270 | users_xtra.insert(), dict(id=2, name="newname2", login_email="newemail2") 271 | ) 272 | 273 | merge = Merge(users, users_xtra, users.c.id == users_xtra.c.id) 274 | merge.when_matched_then_update().where(users_xtra.c.id != 1).values( 275 | name=users_xtra.c.name 276 | ) 277 | 278 | result = connection.execute(merge) 279 | eq_( 280 | connection.execute(users.select().order_by(users.c.id)).fetchall(), 281 | [(1, "name1", "email1"), (2, "newname2", "email2")], 282 | ) 283 | 284 | def test_when_matched_delete(self, connection): 285 | users = self.tables.users 286 | users_xtra = self.tables.users_xtra 287 | 288 | connection.execute( 289 | users.insert(), dict(id=1, name="name1", login_email="email1") 290 | ) 291 | connection.execute( 292 | users.insert(), dict(id=2, name="name2", login_email="email2") 293 | ) 294 | connection.execute( 295 | users_xtra.insert(), dict(id=2, name="newname2", login_email="newemail2") 296 | ) 297 | 298 | merge = Merge(users, users_xtra, users.c.id == users_xtra.c.id) 299 | merge.when_matched_then_delete() 300 | 301 | result = connection.execute(merge) 302 | eq_( 303 | connection.execute(users.select().order_by(users.c.id)).fetchall(), 304 | [(1, "name1", "email1")], 305 | ) 306 | 307 | def test_mixed_criteria(self, connection): 308 | users = self.tables.users 309 | users_xtra = self.tables.users_xtra 310 | 311 | connection.execute( 312 | users.insert(), dict(id=1, name="name1", login_email="email1") 313 | ) 314 | connection.execute( 315 | users.insert(), dict(id=2, name="name2", login_email="email2") 316 | ) 317 | connection.execute( 318 | users_xtra.insert(), dict(id=1, name="newname1", login_email="newemail1") 319 | ) 320 | connection.execute( 321 | users_xtra.insert(), dict(id=2, name="newname2", login_email="newemail2") 322 | ) 323 | connection.execute( 324 | users_xtra.insert(), dict(id=3, name="newname3", login_email="newemail3") 325 | ) 326 | 327 | merge = Merge(users, users_xtra, users.c.id == users_xtra.c.id) 328 | merge.when_matched_then_update().where(users_xtra.c.id == 1) 329 | merge.when_matched_then_delete() 330 | merge.when_not_matched_then_insert() 331 | 332 | result = connection.execute(merge) 333 | eq_( 334 | connection.execute(users.select().order_by(users.c.id)).fetchall(), 335 | [(1, "newname1", "newemail1"), (3, "newname3", "newemail3")], 336 | ) 337 | 338 | def test_no_matches(self, connection): 339 | users = self.tables.users 340 | users_xtra = self.tables.users_xtra 341 | 342 | connection.execute( 343 | users.insert(), dict(id=1, name="name1", login_email="email1") 344 | ) 345 | connection.execute( 346 | users.insert(), dict(id=2, name="name2", login_email="email2") 347 | ) 348 | connection.execute( 349 | users_xtra.insert(), dict(id=1, name="newname1", login_email="newemail1") 350 | ) 351 | connection.execute( 352 | users_xtra.insert(), dict(id=2, name="newname2", login_email="newemail2") 353 | ) 354 | 355 | merge = Merge(users, users_xtra, users.c.id == users_xtra.c.id) 356 | merge.when_not_matched_then_insert().where(users_xtra.c.id == 99) 357 | 358 | result = connection.execute(merge) 359 | eq_( 360 | connection.execute(users.select().order_by(users.c.id)).fetchall(), 361 | [(1, "name1", "email1"), (2, "name2", "email2")], 362 | ) 363 | 364 | # 365 | # def test_selectable_source_when_not_matched_insert(self, connection): 366 | # users = self.tables.users 367 | # users_xtra = self.tables.users_xtra 368 | # 369 | # connection.execute(users.insert(), dict(id=1, name="name1", login_email="email1")) 370 | # connection.execute(users_xtra.insert(), dict(id=1, name="newname1", login_email="newemail1")) 371 | # connection.execute(users_xtra.insert(), dict(id=2, name="newname2", login_email="newemail2")) 372 | # 373 | # merge = Merge(users, users_xtra.select(), users.c.id == users_xtra.c.id) 374 | # merge.when_not_matched_then_insert() 375 | # 376 | # result = connection.execute(merge) 377 | # eq_( 378 | # connection.execute( 379 | # users.select().order_by(users.c.id) 380 | # ).fetchall(), 381 | # [(1, "name1", "email1"), (2, "newname2", "newemail2")], 382 | # ) 383 | -------------------------------------------------------------------------------- /Pipfile.lock: -------------------------------------------------------------------------------- 1 | { 2 | "_meta": { 3 | "hash": { 4 | "sha256": "67cff3649440470951c241bf14d15e31c6a99e02ba0dc2c7dd07e36431964693" 5 | }, 6 | "pipfile-spec": 6, 7 | "requires": { 8 | "python_version": "3.11" 9 | }, 10 | "sources": [ 11 | { 12 | "name": "pypi", 13 | "url": "https://pypi.org/simple", 14 | "verify_ssl": true 15 | } 16 | ] 17 | }, 18 | "default": { 19 | "databend-driver": { 20 | "hashes": [ 21 | "sha256:1072e8e23440571e1ed8251ab5dc951e6e5740a5803b4fe4dc5489dc267a30c0", 22 | "sha256:1522220653a43581fef8446605a3028029fa44869da6a0ee04d4ee3bbbc6ed0f", 23 | "sha256:65f57989a9a5da821fe1f092d641701ddec4f97c588e20f8e03bfe0b2483f362", 24 | "sha256:663b968a2c3146d5b4ac2e6b3b9ef670d159ec7849e353977029f9f7ad479dfb", 25 | "sha256:bbf06205aee8fc4ae35e0d1a5c0e673cb1b14aa2613ded572dc1320404b9aa2d" 26 | ], 27 | "markers": "python_version < '3.14' and python_version >= '3.8'", 28 | "version": "==0.26.2" 29 | }, 30 | "databend-sqlalchemy": { 31 | "editable": true, 32 | "file": ".", 33 | "markers": "python_version >= '3.7'" 34 | }, 35 | "greenlet": { 36 | "hashes": [ 37 | "sha256:0153404a4bb921f0ff1abeb5ce8a5131da56b953eda6e14b88dc6bbc04d2049e", 38 | "sha256:03a088b9de532cbfe2ba2034b2b85e82df37874681e8c470d6fb2f8c04d7e4b7", 39 | "sha256:04b013dc07c96f83134b1e99888e7a79979f1a247e2a9f59697fa14b5862ed01", 40 | "sha256:05175c27cb459dcfc05d026c4232f9de8913ed006d42713cb8a5137bd49375f1", 41 | "sha256:09fc016b73c94e98e29af67ab7b9a879c307c6731a2c9da0db5a7d9b7edd1159", 42 | "sha256:0bbae94a29c9e5c7e4a2b7f0aae5c17e8e90acbfd3bf6270eeba60c39fce3563", 43 | "sha256:0fde093fb93f35ca72a556cf72c92ea3ebfda3d79fc35bb19fbe685853869a83", 44 | "sha256:1443279c19fca463fc33e65ef2a935a5b09bb90f978beab37729e1c3c6c25fe9", 45 | "sha256:1776fd7f989fc6b8d8c8cb8da1f6b82c5814957264d1f6cf818d475ec2bf6395", 46 | "sha256:1d3755bcb2e02de341c55b4fca7a745a24a9e7212ac953f6b3a48d117d7257aa", 47 | "sha256:23f20bb60ae298d7d8656c6ec6db134bca379ecefadb0b19ce6f19d1f232a942", 48 | "sha256:275f72decf9932639c1c6dd1013a1bc266438eb32710016a1c742df5da6e60a1", 49 | "sha256:2846930c65b47d70b9d178e89c7e1a69c95c1f68ea5aa0a58646b7a96df12441", 50 | "sha256:3319aa75e0e0639bc15ff54ca327e8dc7a6fe404003496e3c6925cd3142e0e22", 51 | "sha256:346bed03fe47414091be4ad44786d1bd8bef0c3fcad6ed3dee074a032ab408a9", 52 | "sha256:36b89d13c49216cadb828db8dfa6ce86bbbc476a82d3a6c397f0efae0525bdd0", 53 | "sha256:37b9de5a96111fc15418819ab4c4432e4f3c2ede61e660b1e33971eba26ef9ba", 54 | "sha256:396979749bd95f018296af156201d6211240e7a23090f50a8d5d18c370084dc3", 55 | "sha256:3b2813dc3de8c1ee3f924e4d4227999285fd335d1bcc0d2be6dc3f1f6a318ec1", 56 | "sha256:411f015496fec93c1c8cd4e5238da364e1da7a124bcb293f085bf2860c32c6f6", 57 | "sha256:47da355d8687fd65240c364c90a31569a133b7b60de111c255ef5b606f2ae291", 58 | "sha256:48ca08c771c268a768087b408658e216133aecd835c0ded47ce955381105ba39", 59 | "sha256:4afe7ea89de619adc868e087b4d2359282058479d7cfb94970adf4b55284574d", 60 | "sha256:4ce3ac6cdb6adf7946475d7ef31777c26d94bccc377e070a7986bd2d5c515467", 61 | "sha256:4ead44c85f8ab905852d3de8d86f6f8baf77109f9da589cb4fa142bd3b57b475", 62 | "sha256:54558ea205654b50c438029505def3834e80f0869a70fb15b871c29b4575ddef", 63 | "sha256:5e06afd14cbaf9e00899fae69b24a32f2196c19de08fcb9f4779dd4f004e5e7c", 64 | "sha256:62ee94988d6b4722ce0028644418d93a52429e977d742ca2ccbe1c4f4a792511", 65 | "sha256:63e4844797b975b9af3a3fb8f7866ff08775f5426925e1e0bbcfe7932059a12c", 66 | "sha256:6510bf84a6b643dabba74d3049ead221257603a253d0a9873f55f6a59a65f822", 67 | "sha256:667a9706c970cb552ede35aee17339a18e8f2a87a51fba2ed39ceeeb1004798a", 68 | "sha256:6ef9ea3f137e5711f0dbe5f9263e8c009b7069d8a1acea822bd5e9dae0ae49c8", 69 | "sha256:7017b2be767b9d43cc31416aba48aab0d2309ee31b4dbf10a1d38fb7972bdf9d", 70 | "sha256:7124e16b4c55d417577c2077be379514321916d5790fa287c9ed6f23bd2ffd01", 71 | "sha256:73aaad12ac0ff500f62cebed98d8789198ea0e6f233421059fa68a5aa7220145", 72 | "sha256:77c386de38a60d1dfb8e55b8c1101d68c79dfdd25c7095d51fec2dd800892b80", 73 | "sha256:7876452af029456b3f3549b696bb36a06db7c90747740c5302f74a9e9fa14b13", 74 | "sha256:7939aa3ca7d2a1593596e7ac6d59391ff30281ef280d8632fa03d81f7c5f955e", 75 | "sha256:8320f64b777d00dd7ccdade271eaf0cad6636343293a25074cc5566160e4de7b", 76 | "sha256:85f3ff71e2e60bd4b4932a043fbbe0f499e263c628390b285cb599154a3b03b1", 77 | "sha256:8b8b36671f10ba80e159378df9c4f15c14098c4fd73a36b9ad715f057272fbef", 78 | "sha256:93147c513fac16385d1036b7e5b102c7fbbdb163d556b791f0f11eada7ba65dc", 79 | "sha256:935e943ec47c4afab8965954bf49bfa639c05d4ccf9ef6e924188f762145c0ff", 80 | "sha256:94b6150a85e1b33b40b1464a3f9988dcc5251d6ed06842abff82e42632fac120", 81 | "sha256:94ebba31df2aa506d7b14866fed00ac141a867e63143fe5bca82a8e503b36437", 82 | "sha256:95ffcf719966dd7c453f908e208e14cde192e09fde6c7186c8f1896ef778d8cd", 83 | "sha256:98884ecf2ffb7d7fe6bd517e8eb99d31ff7855a840fa6d0d63cd07c037f6a981", 84 | "sha256:99cfaa2110534e2cf3ba31a7abcac9d328d1d9f1b95beede58294a60348fba36", 85 | "sha256:9e8f8c9cb53cdac7ba9793c276acd90168f416b9ce36799b9b885790f8ad6c0a", 86 | "sha256:a0dfc6c143b519113354e780a50381508139b07d2177cb6ad6a08278ec655798", 87 | "sha256:b2795058c23988728eec1f36a4e5e4ebad22f8320c85f3587b539b9ac84128d7", 88 | "sha256:b42703b1cf69f2aa1df7d1030b9d77d3e584a70755674d60e710f0af570f3761", 89 | "sha256:b7cede291382a78f7bb5f04a529cb18e068dd29e0fb27376074b6d0317bf4dd0", 90 | "sha256:b8a678974d1f3aa55f6cc34dc480169d58f2e6d8958895d68845fa4ab566509e", 91 | "sha256:b8da394b34370874b4572676f36acabac172602abf054cbc4ac910219f3340af", 92 | "sha256:c3a701fe5a9695b238503ce5bbe8218e03c3bcccf7e204e455e7462d770268aa", 93 | "sha256:c4aab7f6381f38a4b42f269057aee279ab0fc7bf2e929e3d4abfae97b682a12c", 94 | "sha256:ca9d0ff5ad43e785350894d97e13633a66e2b50000e8a183a50a88d834752d42", 95 | "sha256:d0028e725ee18175c6e422797c407874da24381ce0690d6b9396c204c7f7276e", 96 | "sha256:d21e10da6ec19b457b82636209cbe2331ff4306b54d06fa04b7c138ba18c8a81", 97 | "sha256:d5e975ca70269d66d17dd995dafc06f1b06e8cb1ec1e9ed54c1d1e4a7c4cf26e", 98 | "sha256:da7a9bff22ce038e19bf62c4dd1ec8391062878710ded0a845bcf47cc0200617", 99 | "sha256:db32b5348615a04b82240cc67983cb315309e88d444a288934ee6ceaebcad6cc", 100 | "sha256:dcc62f31eae24de7f8dce72134c8651c58000d3b1868e01392baea7c32c247de", 101 | "sha256:dfc59d69fc48664bc693842bd57acfdd490acafda1ab52c7836e3fc75c90a111", 102 | "sha256:e347b3bfcf985a05e8c0b7d462ba6f15b1ee1c909e2dcad795e49e91b152c383", 103 | "sha256:e4d333e558953648ca09d64f13e6d8f0523fa705f51cae3f03b5983489958c70", 104 | "sha256:ed10eac5830befbdd0c32f83e8aa6288361597550ba669b04c48f0f9a2c843c6", 105 | "sha256:efc0f674aa41b92da8c49e0346318c6075d734994c3c4e4430b1c3f853e498e4", 106 | "sha256:f1695e76146579f8c06c1509c7ce4dfe0706f49c6831a817ac04eebb2fd02011", 107 | "sha256:f1d4aeb8891338e60d1ab6127af1fe45def5259def8094b9c7e34690c8858803", 108 | "sha256:f406b22b7c9a9b4f8aa9d2ab13d6ae0ac3e85c9a809bd590ad53fed2bf70dc79", 109 | "sha256:f6ff3b14f2df4c41660a7dec01045a045653998784bf8cfcb5a525bdffffbc8f" 110 | ], 111 | "markers": "python_version < '3.14' and platform_machine == 'aarch64' or (platform_machine == 'ppc64le' or (platform_machine == 'x86_64' or (platform_machine == 'amd64' or (platform_machine == 'AMD64' or (platform_machine == 'win32' or platform_machine == 'WIN32')))))", 112 | "version": "==3.1.1" 113 | }, 114 | "sqlalchemy": { 115 | "hashes": [ 116 | "sha256:018ee97c558b499b58935c5a152aeabf6d36b3d55d91656abeb6d93d663c0c4c", 117 | "sha256:01da15490c9df352fbc29859d3c7ba9cd1377791faeeb47c100832004c99472c", 118 | "sha256:04545042969833cb92e13b0a3019549d284fd2423f318b6ba10e7aa687690a3c", 119 | "sha256:06205eb98cb3dd52133ca6818bf5542397f1dd1b69f7ea28aa84413897380b06", 120 | "sha256:08cf721bbd4391a0e765fe0fe8816e81d9f43cece54fdb5ac465c56efafecb3d", 121 | "sha256:0d7e3866eb52d914aea50c9be74184a0feb86f9af8aaaa4daefe52b69378db0b", 122 | "sha256:125a7763b263218a80759ad9ae2f3610aaf2c2fbbd78fff088d584edf81f3782", 123 | "sha256:23c5aa33c01bd898f879db158537d7e7568b503b15aad60ea0c8da8109adf3e7", 124 | "sha256:2600a50d590c22d99c424c394236899ba72f849a02b10e65b4c70149606408b5", 125 | "sha256:2d7332868ce891eda48896131991f7f2be572d65b41a4050957242f8e935d5d7", 126 | "sha256:2ed107331d188a286611cea9022de0afc437dd2d3c168e368169f27aa0f61338", 127 | "sha256:3395e7ed89c6d264d38bea3bfb22ffe868f906a7985d03546ec7dc30221ea980", 128 | "sha256:344cd1ec2b3c6bdd5dfde7ba7e3b879e0f8dd44181f16b895940be9b842fd2b6", 129 | "sha256:34d5c49f18778a3665d707e6286545a30339ad545950773d43977e504815fa70", 130 | "sha256:35e72518615aa5384ef4fae828e3af1b43102458b74a8c481f69af8abf7e802a", 131 | "sha256:3eb14ba1a9d07c88669b7faf8f589be67871d6409305e73e036321d89f1d904e", 132 | "sha256:412c6c126369ddae171c13987b38df5122cb92015cba6f9ee1193b867f3f1530", 133 | "sha256:4600c7a659d381146e1160235918826c50c80994e07c5b26946a3e7ec6c99249", 134 | "sha256:463ecfb907b256e94bfe7bcb31a6d8c7bc96eca7cbe39803e448a58bb9fcad02", 135 | "sha256:4a06e6c8e31c98ddc770734c63903e39f1947c9e3e5e4bef515c5491b7737dde", 136 | "sha256:4b2de1523d46e7016afc7e42db239bd41f2163316935de7c84d0e19af7e69538", 137 | "sha256:4dabd775fd66cf17f31f8625fc0e4cfc5765f7982f94dc09b9e5868182cb71c0", 138 | "sha256:4eff9c270afd23e2746e921e80182872058a7a592017b2713f33f96cc5f82e32", 139 | "sha256:52607d0ebea43cf214e2ee84a6a76bc774176f97c5a774ce33277514875a718e", 140 | "sha256:533e0f66c32093a987a30df3ad6ed21170db9d581d0b38e71396c49718fbb1ca", 141 | "sha256:5493a8120d6fc185f60e7254fc056a6742f1db68c0f849cfc9ab46163c21df47", 142 | "sha256:5d2d1fe548def3267b4c70a8568f108d1fed7cbbeccb9cc166e05af2abc25c22", 143 | "sha256:5dfbc543578058c340360f851ddcecd7a1e26b0d9b5b69259b526da9edfa8875", 144 | "sha256:66a40003bc244e4ad86b72abb9965d304726d05a939e8c09ce844d27af9e6d37", 145 | "sha256:67de057fbcb04a066171bd9ee6bcb58738d89378ee3cabff0bffbf343ae1c787", 146 | "sha256:6827f8c1b2f13f1420545bd6d5b3f9e0b85fe750388425be53d23c760dcf176b", 147 | "sha256:6b35e07f1d57b79b86a7de8ecdcefb78485dab9851b9638c2c793c50203b2ae8", 148 | "sha256:7399d45b62d755e9ebba94eb89437f80512c08edde8c63716552a3aade61eb42", 149 | "sha256:788b6ff6728072b313802be13e88113c33696a9a1f2f6d634a97c20f7ef5ccce", 150 | "sha256:78f1b79132a69fe8bd6b5d91ef433c8eb40688ba782b26f8c9f3d2d9ca23626f", 151 | "sha256:79f4f502125a41b1b3b34449e747a6abfd52a709d539ea7769101696bdca6716", 152 | "sha256:7a8517b6d4005facdbd7eb4e8cf54797dbca100a7df459fdaff4c5123265c1cd", 153 | "sha256:7bd5c5ee1448b6408734eaa29c0d820d061ae18cb17232ce37848376dcfa3e92", 154 | "sha256:7f5243357e6da9a90c56282f64b50d29cba2ee1f745381174caacc50d501b109", 155 | "sha256:805cb481474e111ee3687c9047c5f3286e62496f09c0e82e8853338aaaa348f8", 156 | "sha256:871f55e478b5a648c08dd24af44345406d0e636ffe021d64c9b57a4a11518304", 157 | "sha256:87a1ce1f5e5dc4b6f4e0aac34e7bb535cb23bd4f5d9c799ed1633b65c2bcad8c", 158 | "sha256:8a10ca7f8a1ea0fd5630f02feb055b0f5cdfcd07bb3715fc1b6f8cb72bf114e4", 159 | "sha256:995c2bacdddcb640c2ca558e6760383dcdd68830160af92b5c6e6928ffd259b4", 160 | "sha256:9f03143f8f851dd8de6b0c10784363712058f38209e926723c80654c1b40327a", 161 | "sha256:a1c6b0a5e3e326a466d809b651c63f278b1256146a377a528b6938a279da334f", 162 | "sha256:a28f9c238f1e143ff42ab3ba27990dfb964e5d413c0eb001b88794c5c4a528a9", 163 | "sha256:b2cf5b5ddb69142511d5559c427ff00ec8c0919a1e6c09486e9c32636ea2b9dd", 164 | "sha256:b761a6847f96fdc2d002e29e9e9ac2439c13b919adfd64e8ef49e75f6355c548", 165 | "sha256:bf555f3e25ac3a70c67807b2949bfe15f377a40df84b71ab2c58d8593a1e036e", 166 | "sha256:c08a972cbac2a14810463aec3a47ff218bb00c1a607e6689b531a7c589c50723", 167 | "sha256:c457a38351fb6234781d054260c60e531047e4d07beca1889b558ff73dc2014b", 168 | "sha256:c4c433f78c2908ae352848f56589c02b982d0e741b7905228fad628999799de4", 169 | "sha256:d9f119e7736967c0ea03aff91ac7d04555ee038caf89bb855d93bbd04ae85b41", 170 | "sha256:e6b0a1c7ed54a5361aaebb910c1fa864bae34273662bb4ff788a527eafd6e14d", 171 | "sha256:f2bcb085faffcacf9319b1b1445a7e1cfdc6fb46c03f2dce7bc2d9a4b3c1cdc5", 172 | "sha256:fe193d3ae297c423e0e567e240b4324d6b6c280a048e64c77a3ea6886cc2aa87" 173 | ], 174 | "markers": "python_version >= '3.7'", 175 | "version": "==2.0.39" 176 | }, 177 | "typing-extensions": { 178 | "hashes": [ 179 | "sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d", 180 | "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8" 181 | ], 182 | "markers": "python_version >= '3.8'", 183 | "version": "==4.12.2" 184 | } 185 | }, 186 | "develop": { 187 | "black": { 188 | "hashes": [ 189 | "sha256:030b9759066a4ee5e5aca28c3c77f9c64789cdd4de8ac1df642c40b708be6171", 190 | "sha256:055e59b198df7ac0b7efca5ad7ff2516bca343276c466be72eb04a3bcc1f82d7", 191 | "sha256:0e519ecf93120f34243e6b0054db49c00a35f84f195d5bce7e9f5cfc578fc2da", 192 | "sha256:172b1dbff09f86ce6f4eb8edf9dede08b1fce58ba194c87d7a4f1a5aa2f5b3c2", 193 | "sha256:1e2978f6df243b155ef5fa7e558a43037c3079093ed5d10fd84c43900f2d8ecc", 194 | "sha256:33496d5cd1222ad73391352b4ae8da15253c5de89b93a80b3e2c8d9a19ec2666", 195 | "sha256:3b48735872ec535027d979e8dcb20bf4f70b5ac75a8ea99f127c106a7d7aba9f", 196 | "sha256:4b60580e829091e6f9238c848ea6750efed72140b91b048770b64e74fe04908b", 197 | "sha256:759e7ec1e050a15f89b770cefbf91ebee8917aac5c20483bc2d80a6c3a04df32", 198 | "sha256:8f0b18a02996a836cc9c9c78e5babec10930862827b1b724ddfe98ccf2f2fe4f", 199 | "sha256:95e8176dae143ba9097f351d174fdaf0ccd29efb414b362ae3fd72bf0f710717", 200 | "sha256:96c1c7cd856bba8e20094e36e0f948718dc688dba4a9d78c3adde52b9e6c2299", 201 | "sha256:a1ee0a0c330f7b5130ce0caed9936a904793576ef4d2b98c40835d6a65afa6a0", 202 | "sha256:a22f402b410566e2d1c950708c77ebf5ebd5d0d88a6a2e87c86d9fb48afa0d18", 203 | "sha256:a39337598244de4bae26475f77dda852ea00a93bd4c728e09eacd827ec929df0", 204 | "sha256:afebb7098bfbc70037a053b91ae8437c3857482d3a690fefc03e9ff7aa9a5fd3", 205 | "sha256:bacabb307dca5ebaf9c118d2d2f6903da0d62c9faa82bd21a33eecc319559355", 206 | "sha256:bce2e264d59c91e52d8000d507eb20a9aca4a778731a08cfff7e5ac4a4bb7096", 207 | "sha256:d9e6827d563a2c820772b32ce8a42828dc6790f095f441beef18f96aa6f8294e", 208 | "sha256:db8ea9917d6f8fc62abd90d944920d95e73c83a5ee3383493e35d271aca872e9", 209 | "sha256:ea0213189960bda9cf99be5b8c8ce66bb054af5e9e861249cd23471bd7b0b3ba", 210 | "sha256:f3df5f1bf91d36002b0a75389ca8663510cf0531cca8aa5c1ef695b46d98655f" 211 | ], 212 | "index": "pypi", 213 | "markers": "python_version >= '3.9'", 214 | "version": "==25.1.0" 215 | }, 216 | "click": { 217 | "hashes": [ 218 | "sha256:63c132bbbed01578a06712a2d1f497bb62d9c1c0d329b7903a866228027263b2", 219 | "sha256:ed53c9d8990d83c2a27deae68e4ee337473f6330c040a31d4225c9574d16096a" 220 | ], 221 | "markers": "python_version >= '3.7'", 222 | "version": "==8.1.8" 223 | }, 224 | "execnet": { 225 | "hashes": [ 226 | "sha256:26dee51f1b80cebd6d0ca8e74dd8745419761d3bef34163928cbebbdc4749fdc", 227 | "sha256:5189b52c6121c24feae288166ab41b32549c7e2348652736540b9e6e7d4e72e3" 228 | ], 229 | "markers": "python_version >= '3.8'", 230 | "version": "==2.1.1" 231 | }, 232 | "flake8": { 233 | "hashes": [ 234 | "sha256:1cbc62e65536f65e6d754dfe6f1bada7f5cf392d6f5db3c2b85892466c3e7c1a", 235 | "sha256:c586ffd0b41540951ae41af572e6790dbd49fc12b3aa2541685d253d9bd504bd" 236 | ], 237 | "index": "pypi", 238 | "markers": "python_full_version >= '3.8.1'", 239 | "version": "==7.1.2" 240 | }, 241 | "iniconfig": { 242 | "hashes": [ 243 | "sha256:2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3", 244 | "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374" 245 | ], 246 | "markers": "python_version >= '3.7'", 247 | "version": "==2.0.0" 248 | }, 249 | "mccabe": { 250 | "hashes": [ 251 | "sha256:348e0240c33b60bbdf4e523192ef919f28cb2c3d7d5c7794f74009290f236325", 252 | "sha256:6c2d30ab6be0e4a46919781807b4f0d834ebdd6c6e3dca0bda5a15f863427b6e" 253 | ], 254 | "markers": "python_version >= '3.6'", 255 | "version": "==0.7.0" 256 | }, 257 | "mypy-extensions": { 258 | "hashes": [ 259 | "sha256:4392f6c0eb8a5668a69e23d168ffa70f0be9ccfd32b5cc2d26a34ae5b844552d", 260 | "sha256:75dbf8955dc00442a438fc4d0666508a9a97b6bd41aa2f0ffe9d2f2725af0782" 261 | ], 262 | "markers": "python_version >= '3.5'", 263 | "version": "==1.0.0" 264 | }, 265 | "packaging": { 266 | "hashes": [ 267 | "sha256:09abb1bccd265c01f4a3aa3f7a7db064b36514d2cba19a2f694fe6150451a759", 268 | "sha256:c228a6dc5e932d346bc5739379109d49e8853dd8223571c7c5b55260edc0b97f" 269 | ], 270 | "markers": "python_version >= '3.8'", 271 | "version": "==24.2" 272 | }, 273 | "pathspec": { 274 | "hashes": [ 275 | "sha256:a0d503e138a4c123b27490a4f7beda6a01c6f288df0e4a8b79c7eb0dc7b4cc08", 276 | "sha256:a482d51503a1ab33b1c67a6c3813a26953dbdc71c31dacaef9a838c4e29f5712" 277 | ], 278 | "markers": "python_version >= '3.8'", 279 | "version": "==0.12.1" 280 | }, 281 | "platformdirs": { 282 | "hashes": [ 283 | "sha256:357fb2acbc885b0419afd3ce3ed34564c13c9b95c89360cd9563f73aa5e2b907", 284 | "sha256:73e575e1408ab8103900836b97580d5307456908a03e92031bab39e4554cc3fb" 285 | ], 286 | "markers": "python_version >= '3.8'", 287 | "version": "==4.3.6" 288 | }, 289 | "pluggy": { 290 | "hashes": [ 291 | "sha256:2cffa88e94fdc978c4c574f15f9e59b7f4201d439195c3715ca9e2486f1d0cf1", 292 | "sha256:44e1ad92c8ca002de6377e165f3e0f1be63266ab4d554740532335b9d75ea669" 293 | ], 294 | "markers": "python_version >= '3.8'", 295 | "version": "==1.5.0" 296 | }, 297 | "pycodestyle": { 298 | "hashes": [ 299 | "sha256:46f0fb92069a7c28ab7bb558f05bfc0110dac69a0cd23c61ea0040283a9d78b3", 300 | "sha256:6838eae08bbce4f6accd5d5572075c63626a15ee3e6f842df996bf62f6d73521" 301 | ], 302 | "markers": "python_version >= '3.8'", 303 | "version": "==2.12.1" 304 | }, 305 | "pyflakes": { 306 | "hashes": [ 307 | "sha256:1c61603ff154621fb2a9172037d84dca3500def8c8b630657d1701f026f8af3f", 308 | "sha256:84b5be138a2dfbb40689ca07e2152deb896a65c3a3e24c251c5c62489568074a" 309 | ], 310 | "markers": "python_version >= '3.8'", 311 | "version": "==3.2.0" 312 | }, 313 | "pytest": { 314 | "hashes": [ 315 | "sha256:c69214aa47deac29fad6c2a4f590b9c4a9fdb16a403176fe154b79c0b4d4d820", 316 | "sha256:f4efe70cc14e511565ac476b57c279e12a855b11f48f212af1080ef2263d3845" 317 | ], 318 | "index": "pypi", 319 | "markers": "python_version >= '3.8'", 320 | "version": "==8.3.5" 321 | }, 322 | "pytest-xdist": { 323 | "hashes": [ 324 | "sha256:9ed4adfb68a016610848639bb7e02c9352d5d9f03d04809919e2dafc3be4cca7", 325 | "sha256:ead156a4db231eec769737f57668ef58a2084a34b2e55c4a8fa20d861107300d" 326 | ], 327 | "index": "pypi", 328 | "markers": "python_version >= '3.8'", 329 | "version": "==3.6.1" 330 | } 331 | } 332 | } 333 | -------------------------------------------------------------------------------- /databend_sqlalchemy/dml.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # 3 | # Note: parts of the file come from https://github.com/snowflakedb/snowflake-sqlalchemy 4 | # licensed under the same Apache 2.0 License 5 | from enum import Enum 6 | from types import NoneType 7 | from urllib.parse import urlparse 8 | 9 | from sqlalchemy.sql.selectable import Select, Subquery, TableClause 10 | from sqlalchemy.sql.dml import UpdateBase 11 | from sqlalchemy.sql.elements import ClauseElement 12 | from sqlalchemy.sql.expression import select 13 | from sqlalchemy.sql.roles import FromClauseRole 14 | 15 | 16 | class _OnMergeBaseClause(ClauseElement): 17 | # __visit_name__ = "on_merge_base_clause" 18 | 19 | def __init__(self): 20 | self.set = {} 21 | self.predicate = None 22 | 23 | def __repr__(self): 24 | return f" AND {str(self.predicate)}" if self.predicate is not None else "" 25 | 26 | def values(self, **kwargs): 27 | self.set = kwargs 28 | return self 29 | 30 | def where(self, expr): 31 | self.predicate = expr 32 | return self 33 | 34 | 35 | class WhenMergeMatchedUpdateClause(_OnMergeBaseClause): 36 | __visit_name__ = "when_merge_matched_update" 37 | 38 | def __repr__(self): 39 | case_predicate = super() 40 | update_str = f"WHEN MATCHED{case_predicate} THEN UPDATE" 41 | if not self.set: 42 | return f"{update_str} *" 43 | 44 | set_values = ", ".join( 45 | [f"{set_item[0]} = {set_item[1]}" for set_item in self.set.items()] 46 | ) 47 | return f"{update_str} SET {str(set_values)}" 48 | 49 | 50 | class WhenMergeMatchedDeleteClause(_OnMergeBaseClause): 51 | __visit_name__ = "when_merge_matched_delete" 52 | 53 | def __repr__(self): 54 | case_predicate = super() 55 | return f"WHEN MATCHED{case_predicate} THEN DELETE" 56 | 57 | 58 | class WhenMergeUnMatchedClause(_OnMergeBaseClause): 59 | __visit_name__ = "when_merge_unmatched" 60 | 61 | def __repr__(self): 62 | case_predicate = super() 63 | insert_str = f"WHEN NOT MATCHED{case_predicate} THEN INSERT" 64 | if not self.set: 65 | return f"{insert_str} *" 66 | 67 | sets, sets_tos = zip(*self.set.items()) 68 | return "{} ({}) VALUES ({})".format( 69 | insert_str, 70 | ", ".join(sets), 71 | ", ".join(map(str, sets_tos)), 72 | ) 73 | 74 | 75 | class Merge(UpdateBase): 76 | __visit_name__ = "merge" 77 | _bind = None 78 | 79 | inherit_cache = False 80 | 81 | def __init__(self, target, source, on): 82 | if not isinstance(source, (TableClause, Select, Subquery)): 83 | raise Exception(f"Invalid type for merge source: {source}") 84 | self.target = target 85 | self.source = source 86 | self.on = on 87 | self.clauses = [] 88 | 89 | def __repr__(self): 90 | clauses = " ".join([repr(clause) for clause in self.clauses]) 91 | return ( 92 | f"MERGE INTO {self.target} USING ({select(self.source)}) AS {self.source.name} ON {self.on}" 93 | + (f" {clauses}" if clauses else "") 94 | ) 95 | 96 | def when_matched_then_update(self): 97 | clause = WhenMergeMatchedUpdateClause() 98 | self.clauses.append(clause) 99 | return clause 100 | 101 | def when_matched_then_delete(self): 102 | clause = WhenMergeMatchedDeleteClause() 103 | self.clauses.append(clause) 104 | return clause 105 | 106 | def when_not_matched_then_insert(self): 107 | clause = WhenMergeUnMatchedClause() 108 | self.clauses.append(clause) 109 | return clause 110 | 111 | 112 | class _CopyIntoBase(UpdateBase): 113 | __visit_name__ = "copy_into" 114 | _bind = None 115 | 116 | def __init__( 117 | self, 118 | target: ["TableClause", "StageClause", "_StorageClause"], 119 | from_, 120 | file_format: "CopyFormat" = None, 121 | options: ["CopyIntoLocationOptions", "CopyIntoTableOptions"] = None, 122 | ): 123 | self.target = target 124 | self.from_ = from_ 125 | self.file_format = file_format 126 | self.options = options 127 | 128 | def __repr__(self): 129 | """ 130 | repr for debugging / logging purposes only. For compilation logic, see 131 | the corresponding visitor in base.py 132 | """ 133 | val = f"COPY INTO {self.target} FROM {repr(self.from_)}" 134 | return val + f" {repr(self.file_format)} ({self.options})" 135 | 136 | def bind(self): 137 | return None 138 | 139 | 140 | class CopyIntoLocation(_CopyIntoBase): 141 | inherit_cache = False 142 | 143 | def __init__( 144 | self, 145 | *, 146 | target: ["StageClause", "_StorageClause"], 147 | from_, 148 | file_format: "CopyFormat" = None, 149 | options: "CopyIntoLocationOptions" = None, 150 | ): 151 | super().__init__(target, from_, file_format, options) 152 | 153 | 154 | class CopyIntoTable(_CopyIntoBase): 155 | inherit_cache = False 156 | 157 | def __init__( 158 | self, 159 | *, 160 | target: [TableClause], 161 | from_: ["StageClause", "_StorageClause", "FileColumnClause"], 162 | files: list = None, 163 | pattern: str = None, 164 | file_format: "CopyFormat" = None, 165 | options: "CopyIntoTableOptions" = None, 166 | ): 167 | super().__init__(target, from_, file_format, options) 168 | self.files = files 169 | self.pattern = pattern 170 | 171 | 172 | class _CopyIntoOptions(ClauseElement): 173 | __visit_name__ = "copy_into_options" 174 | 175 | def __init__(self): 176 | self.options = dict() 177 | 178 | def __repr__(self): 179 | return "\n".join([f"{k} = {v}" for k, v in self.options.items()]) 180 | 181 | 182 | class CopyIntoLocationOptions(_CopyIntoOptions): 183 | # __visit_name__ = "copy_into_location_options" 184 | 185 | def __init__( 186 | self, 187 | *, 188 | single: bool = None, 189 | max_file_size_bytes: int = None, 190 | overwrite: bool = None, 191 | include_query_id: bool = None, 192 | use_raw_path: bool = None, 193 | ): 194 | super().__init__() 195 | if not isinstance(single, NoneType): 196 | self.options["SINGLE"] = "TRUE" if single else "FALSE" 197 | if not isinstance(max_file_size_bytes, NoneType): 198 | self.options["MAX_FILE_SIZE "] = max_file_size_bytes 199 | if not isinstance(overwrite, NoneType): 200 | self.options["OVERWRITE"] = "TRUE" if overwrite else "FALSE" 201 | if not isinstance(include_query_id, NoneType): 202 | self.options["INCLUDE_QUERY_ID"] = "TRUE" if include_query_id else "FALSE" 203 | if not isinstance(use_raw_path, NoneType): 204 | self.options["USE_RAW_PATH"] = "TRUE" if use_raw_path else "FALSE" 205 | 206 | 207 | class CopyIntoTableOptions(_CopyIntoOptions): 208 | # __visit_name__ = "copy_into_table_options" 209 | 210 | def __init__( 211 | self, 212 | *, 213 | size_limit: int = None, 214 | purge: bool = None, 215 | force: bool = None, 216 | disable_variant_check: bool = None, 217 | on_error: str = None, 218 | max_files: int = None, 219 | return_failed_only: bool = None, 220 | column_match_mode: str = None, 221 | ): 222 | super().__init__() 223 | if not isinstance(size_limit, NoneType): 224 | self.options["SIZE_LIMIT"] = size_limit 225 | if not isinstance(purge, NoneType): 226 | self.options["PURGE "] = "TRUE" if purge else "FALSE" 227 | if not isinstance(force, NoneType): 228 | self.options["FORCE"] = "TRUE" if force else "FALSE" 229 | if not isinstance(disable_variant_check, NoneType): 230 | self.options["DISABLE_VARIANT_CHECK"] = ( 231 | "TRUE" if disable_variant_check else "FALSE" 232 | ) 233 | if not isinstance(on_error, NoneType): 234 | self.options["ON_ERROR"] = on_error 235 | if not isinstance(max_files, NoneType): 236 | self.options["MAX_FILES"] = max_files 237 | if not isinstance(return_failed_only, NoneType): 238 | self.options["RETURN_FAILED_ONLY"] = return_failed_only 239 | if not isinstance(column_match_mode, NoneType): 240 | self.options["COLUMN_MATCH_MODE"] = column_match_mode 241 | 242 | 243 | class Compression(Enum): 244 | NONE = "NONE" 245 | AUTO = "AUTO" 246 | GZIP = "GZIP" 247 | BZ2 = "BZ2" 248 | BROTLI = "BROTLI" 249 | ZSTD = "ZSTD" 250 | DEFLATE = "DEFLATE" 251 | RAW_DEFLATE = "RAW_DEFLATE" 252 | XZ = "XZ" 253 | SNAPPY = "SNAPPY" 254 | ZIP = "ZIP" 255 | 256 | 257 | class CopyFormat(ClauseElement): 258 | """ 259 | Base class for Format specifications inside a COPY INTO statement. May also 260 | be used to create a named format. 261 | """ 262 | 263 | __visit_name__ = "copy_format" 264 | 265 | def __init__(self, format_name=None): 266 | self.options = dict() 267 | if format_name: 268 | self.options["format_name"] = format_name 269 | 270 | def __repr__(self): 271 | """ 272 | repr for debugging / logging purposes only. For compilation logic, see 273 | the respective visitor in the dialect 274 | """ 275 | return f"FILE_FORMAT=({self.options})" 276 | 277 | 278 | class CSVFormat(CopyFormat): 279 | format_type = "CSV" 280 | 281 | def __init__( 282 | self, 283 | *, 284 | record_delimiter: str = None, 285 | field_delimiter: str = None, 286 | quote: str = None, 287 | escape: str = None, 288 | skip_header: int = None, 289 | nan_display: str = None, 290 | null_display: str = None, 291 | error_on_column_mismatch: bool = None, 292 | empty_field_as: str = None, 293 | output_header: bool = None, 294 | binary_format: str = None, 295 | compression: Compression = None, 296 | ): 297 | super().__init__() 298 | if record_delimiter: 299 | if ( 300 | len(str(record_delimiter).encode().decode("unicode_escape")) != 1 301 | and record_delimiter != "\r\n" 302 | ): 303 | raise TypeError("Record Delimiter should be a single character.") 304 | self.options["RECORD_DELIMITER"] = f"{repr(record_delimiter)}" 305 | if field_delimiter: 306 | if len(str(field_delimiter).encode().decode("unicode_escape")) != 1: 307 | raise TypeError("Field Delimiter should be a single character") 308 | self.options["FIELD_DELIMITER"] = f"{repr(field_delimiter)}" 309 | if quote: 310 | if quote not in ["'", '"', "`"]: 311 | raise TypeError("Quote character must be one of [', \", `].") 312 | self.options["QUOTE"] = f"{repr(quote)}" 313 | if escape: 314 | if escape not in ["\\", ""]: 315 | raise TypeError('Escape character must be "\\" or "".') 316 | self.options["ESCAPE"] = f"{repr(escape)}" 317 | if skip_header: 318 | if skip_header < 0: 319 | raise TypeError("Skip header must be positive integer.") 320 | self.options["SKIP_HEADER"] = skip_header 321 | if nan_display: 322 | if nan_display not in ["NULL", "NaN"]: 323 | raise TypeError('NaN Display should be "NULL" or "NaN".') 324 | self.options["NAN_DISPLAY"] = f"'{nan_display}'" 325 | if null_display: 326 | self.options["NULL_DISPLAY"] = f"'{null_display}'" 327 | if error_on_column_mismatch: 328 | self.options["ERROR_ON_COLUMN_MISMATCH"] = str( 329 | error_on_column_mismatch 330 | ).upper() 331 | if empty_field_as: 332 | if empty_field_as not in ["NULL", "STRING", "FIELD_DEFAULT"]: 333 | raise TypeError( 334 | 'Empty Field As should be "NULL", "STRING" for "FIELD_DEFAULT".' 335 | ) 336 | self.options["EMPTY_FIELD_AS"] = f"{empty_field_as}" 337 | if output_header: 338 | self.options["OUTPUT_HEADER"] = str(output_header).upper() 339 | if binary_format: 340 | if binary_format not in ["HEX", "BASE64"]: 341 | raise TypeError('Binary Format should be "HEX" or "BASE64".') 342 | self.options["BINARY_FORMAT"] = binary_format 343 | if compression: 344 | self.options["COMPRESSION"] = compression.value 345 | 346 | 347 | class TSVFormat(CopyFormat): 348 | format_type = "TSV" 349 | 350 | def __init__( 351 | self, 352 | *, 353 | record_delimiter: str = None, 354 | field_delimiter: str = None, 355 | compression: Compression = None, 356 | ): 357 | super().__init__() 358 | if record_delimiter: 359 | if ( 360 | len(str(record_delimiter).encode().decode("unicode_escape")) != 1 361 | and record_delimiter != "\r\n" 362 | ): 363 | raise TypeError("Record Delimiter should be a single character.") 364 | self.options["RECORD_DELIMITER"] = f"{repr(record_delimiter)}" 365 | if field_delimiter: 366 | if len(str(field_delimiter).encode().decode("unicode_escape")) != 1: 367 | raise TypeError("Field Delimiter should be a single character") 368 | self.options["FIELD_DELIMITER"] = f"{repr(field_delimiter)}" 369 | if compression: 370 | self.options["COMPRESSION"] = compression.value 371 | 372 | 373 | class NDJSONFormat(CopyFormat): 374 | format_type = "NDJSON" 375 | 376 | def __init__( 377 | self, 378 | *, 379 | null_field_as: str = None, 380 | missing_field_as: str = None, 381 | compression: Compression = None, 382 | ): 383 | super().__init__() 384 | if null_field_as: 385 | if null_field_as not in ["NULL", "FIELD_DEFAULT"]: 386 | raise TypeError('Null Field As should be "NULL" or "FIELD_DEFAULT".') 387 | self.options["NULL_FIELD_AS"] = f"{null_field_as}" 388 | if missing_field_as: 389 | if missing_field_as not in [ 390 | "ERROR", 391 | "NULL", 392 | "FIELD_DEFAULT", 393 | "TYPE_DEFAULT", 394 | ]: 395 | raise TypeError( 396 | 'Missing Field As should be "ERROR", "NULL", "FIELD_DEFAULT" or "TYPE_DEFAULT".' 397 | ) 398 | self.options["MISSING_FIELD_AS"] = f"{missing_field_as}" 399 | if compression: 400 | self.options["COMPRESSION"] = compression.value 401 | 402 | 403 | class ParquetFormat(CopyFormat): 404 | format_type = "PARQUET" 405 | 406 | def __init__( 407 | self, 408 | *, 409 | missing_field_as: str = None, 410 | compression: Compression = None, 411 | ): 412 | super().__init__() 413 | if missing_field_as: 414 | if missing_field_as not in ["ERROR", "FIELD_DEFAULT"]: 415 | raise TypeError( 416 | 'Missing Field As should be "ERROR" or "FIELD_DEFAULT".' 417 | ) 418 | self.options["MISSING_FIELD_AS"] = f"{missing_field_as}" 419 | if compression: 420 | if compression not in [Compression.ZSTD, Compression.SNAPPY]: 421 | raise TypeError( 422 | 'Compression should be None, ZStd, or Snappy.' 423 | ) 424 | self.options["COMPRESSION"] = compression.value 425 | 426 | 427 | class AVROFormat(CopyFormat): 428 | format_type = "AVRO" 429 | 430 | def __init__( 431 | self, 432 | *, 433 | missing_field_as: str = None, 434 | ): 435 | super().__init__() 436 | if missing_field_as: 437 | if missing_field_as not in ["ERROR", "FIELD_DEFAULT"]: 438 | raise TypeError( 439 | 'Missing Field As should be "ERROR" or "FIELD_DEFAULT".' 440 | ) 441 | self.options["MISSING_FIELD_AS"] = f"{missing_field_as}" 442 | 443 | 444 | class ORCFormat(CopyFormat): 445 | format_type = "ORC" 446 | 447 | def __init__( 448 | self, 449 | *, 450 | missing_field_as: str = None, 451 | ): 452 | super().__init__() 453 | if missing_field_as: 454 | if missing_field_as not in ["ERROR", "FIELD_DEFAULT"]: 455 | raise TypeError( 456 | 'Missing Field As should be "ERROR" or "FIELD_DEFAULT".' 457 | ) 458 | self.options["MISSING_FIELD_AS"] = f"{missing_field_as}" 459 | 460 | 461 | class StageClause(ClauseElement, FromClauseRole): 462 | """Stage Clause""" 463 | 464 | __visit_name__ = "stage" 465 | _hide_froms = () 466 | 467 | def __init__(self, *, name, path=None): 468 | self.name = name 469 | self.path = path 470 | 471 | def __repr__(self): 472 | return f"@{self.name}/{self.path}" 473 | 474 | 475 | class FileColumnClause(ClauseElement, FromClauseRole): 476 | """Clause for selecting file columns from a Stage/Location""" 477 | 478 | __visit_name__ = "file_column" 479 | 480 | def __init__(self, *, columns, from_: ["StageClause", "_StorageClause"]): 481 | # columns need to be expressions of column index, e.g. $1, IF($1 =='t', True, False), or string of these expressions that we just use 482 | self.columns = columns 483 | self.from_ = from_ 484 | 485 | def __repr__(self): 486 | return ( 487 | f"SELECT {self.columns if isinstance(self.columns, str) else ','.join(repr(col) for col in self.columns)}" 488 | f" FROM {repr(self.from_)}" 489 | ) 490 | 491 | 492 | class _StorageClause(ClauseElement): 493 | pass 494 | 495 | 496 | class AmazonS3(_StorageClause): 497 | """Amazon S3""" 498 | 499 | __visit_name__ = "amazon_s3" 500 | 501 | def __init__( 502 | self, 503 | uri: str, 504 | access_key_id: str, 505 | secret_access_key: str, 506 | endpoint_url: str = None, 507 | enable_virtual_host_style: bool = None, 508 | master_key: str = None, 509 | region: str = None, 510 | security_token: str = None, 511 | ): 512 | r = urlparse(uri) 513 | if r.scheme != "s3": 514 | raise ValueError(f"Invalid S3 URI: {uri}") 515 | 516 | self.uri = uri 517 | self.access_key_id = access_key_id 518 | self.secret_access_key = secret_access_key 519 | self.bucket = r.netloc 520 | self.path = r.path 521 | if endpoint_url: 522 | self.endpoint_url = endpoint_url 523 | if enable_virtual_host_style: 524 | self.enable_virtual_host_style = enable_virtual_host_style 525 | if master_key: 526 | self.master_key = master_key 527 | if region: 528 | self.region = region 529 | if security_token: 530 | self.security_token = security_token 531 | 532 | def __repr__(self): 533 | return ( 534 | f"'{self.uri}' \n" 535 | f"CONNECTION = (\n" 536 | f" ENDPOINT_URL = '{self.endpoint_url}' \n" 537 | if self.endpoint_url 538 | else ( 539 | "" 540 | f" ACCESS_KEY_ID = '{self.access_key_id}' \n" 541 | f" SECRET_ACCESS_KEY = '{self.secret_access_key}'\n" 542 | f" ENABLE_VIRTUAL_HOST_STYLE = '{self.enable_virtual_host_style}'\n" 543 | if self.enable_virtual_host_style 544 | else ( 545 | "" f" MASTER_KEY = '{self.master_key}'\n" 546 | if self.master_key 547 | else ( 548 | "" f" REGION = '{self.region}'\n" 549 | if self.region 550 | else ( 551 | "" f" SECURITY_TOKEN = '{self.security_token}'\n" 552 | if self.security_token 553 | else "" f")" 554 | ) 555 | ) 556 | ) 557 | ) 558 | ) 559 | 560 | 561 | class AzureBlobStorage(_StorageClause): 562 | """Microsoft Azure Blob Storage""" 563 | 564 | __visit_name__ = "azure_blob_storage" 565 | 566 | def __init__(self, *, uri: str, account_name: str, account_key: str): 567 | r = urlparse(uri) 568 | if r.scheme != "azblob": 569 | raise ValueError(f"Invalid Azure URI: {uri}") 570 | 571 | self.uri = uri 572 | self.account_name = account_name 573 | self.account_key = account_key 574 | self.container = r.netloc 575 | self.path = r.path 576 | 577 | def __repr__(self): 578 | return ( 579 | f"'{self.uri}' \n" 580 | f"CONNECTION = (\n" 581 | f" ENDPOINT_URL = 'https://{self.account_name}.blob.core.windows.net' \n" 582 | f" ACCOUNT_NAME = '{self.account_name}' \n" 583 | f" ACCOUNT_KEY = '{self.account_key}'\n" 584 | f")" 585 | ) 586 | 587 | 588 | class GoogleCloudStorage(_StorageClause): 589 | """Google Cloud Storage""" 590 | 591 | __visit_name__ = "google_cloud_storage" 592 | 593 | def __init__(self, *, uri, credentials): 594 | r = urlparse(uri) 595 | if r.scheme != "gcs": 596 | raise ValueError(f"Invalid Google Cloud Storage URI: {uri}") 597 | 598 | self.uri = uri 599 | self.credentials = credentials 600 | 601 | def __repr__(self): 602 | return ( 603 | f"'{self.uri}' \n" 604 | f"CONNECTION = (\n" 605 | f" ENDPOINT_URL = 'https://storage.googleapis.com' \n" 606 | f" CREDENTIAL = '{self.credentials}' \n" 607 | f")" 608 | ) 609 | -------------------------------------------------------------------------------- /tests/test_sqlalchemy.py: -------------------------------------------------------------------------------- 1 | # tests/test_suite.py 2 | 3 | from sqlalchemy.testing.suite import * 4 | 5 | from sqlalchemy.testing.suite import ComponentReflectionTestExtra as _ComponentReflectionTestExtra 6 | from sqlalchemy.testing.suite import DeprecatedCompoundSelectTest as _DeprecatedCompoundSelectTest 7 | from sqlalchemy.testing.suite import BooleanTest as _BooleanTest 8 | from sqlalchemy.testing.suite import BinaryTest as _BinaryTest 9 | from sqlalchemy.testing.suite import CompoundSelectTest as _CompoundSelectTest 10 | from sqlalchemy.testing.suite import HasIndexTest as _HasIndexTest 11 | from sqlalchemy.testing.suite import InsertBehaviorTest as _InsertBehaviorTest 12 | from sqlalchemy.testing.suite import LikeFunctionsTest as _LikeFunctionsTest 13 | from sqlalchemy.testing.suite import LongNameBlowoutTest as _LongNameBlowoutTest 14 | from sqlalchemy.testing.suite import QuotedNameArgumentTest as _QuotedNameArgumentTest 15 | from sqlalchemy.testing.suite import JoinTest as _JoinTest 16 | from sqlalchemy.testing.suite import HasSequenceTest as _HasSequenceTest 17 | 18 | from sqlalchemy.testing.suite import ServerSideCursorsTest as _ServerSideCursorsTest 19 | 20 | from sqlalchemy.testing.suite import CTETest as _CTETest 21 | from sqlalchemy.testing.suite import JSONTest as _JSONTest 22 | from sqlalchemy.testing.suite import IntegerTest as _IntegerTest 23 | 24 | from sqlalchemy import types as sql_types 25 | from sqlalchemy.testing import config, skip_test 26 | from sqlalchemy import testing, Table, Column, Integer 27 | from sqlalchemy.testing import eq_, fixtures, assertions 28 | 29 | from databend_sqlalchemy.types import TINYINT, BITMAP, DOUBLE, GEOMETRY, GEOGRAPHY 30 | 31 | from packaging import version 32 | import sqlalchemy 33 | if version.parse(sqlalchemy.__version__) >= version.parse('2.0.0'): 34 | if version.parse(sqlalchemy.__version__) < version.parse('2.0.42'): 35 | from sqlalchemy.testing.suite import BizarroCharacterFKResolutionTest as _BizarroCharacterFKResolutionTest 36 | from sqlalchemy.testing.suite import EnumTest as _EnumTest 37 | else: 38 | from sqlalchemy.testing.suite import ComponentReflectionTest as _ComponentReflectionTest 39 | 40 | class ComponentReflectionTest(_ComponentReflectionTest): 41 | 42 | @testing.skip("databend") 43 | def test_get_indexes(self): 44 | pass 45 | 46 | class ComponentReflectionTestExtra(_ComponentReflectionTestExtra): 47 | @testing.skip("databend") #ToDo No length in Databend 48 | @testing.requires.table_reflection 49 | def test_varchar_reflection(self, connection, metadata): 50 | typ = self._type_round_trip( 51 | connection, metadata, sql_types.String(52) 52 | )[0] 53 | assert isinstance(typ, sql_types.String) 54 | eq_(typ.length, 52) 55 | 56 | @testing.skip("databend") # ToDo No length in Databend 57 | @testing.requires.table_reflection 58 | @testing.combinations( 59 | sql_types.String, 60 | sql_types.VARCHAR, 61 | sql_types.CHAR, 62 | (sql_types.NVARCHAR, testing.requires.nvarchar_types), 63 | (sql_types.NCHAR, testing.requires.nvarchar_types), 64 | argnames="type_", 65 | ) 66 | def test_string_length_reflection(self, connection, metadata, type_): 67 | typ = self._type_round_trip(connection, metadata, type_(52))[0] 68 | if issubclass(type_, sql_types.VARCHAR): 69 | assert isinstance(typ, sql_types.VARCHAR) 70 | elif issubclass(type_, sql_types.CHAR): 71 | assert isinstance(typ, sql_types.CHAR) 72 | else: 73 | assert isinstance(typ, sql_types.String) 74 | 75 | eq_(typ.length, 52) 76 | assert isinstance(typ.length, int) 77 | 78 | 79 | class BooleanTest(_BooleanTest): 80 | __backend__ = True 81 | 82 | def test_whereclause(self): 83 | """ 84 | This is overridden from Ancestor implementation because Databend does not support `WHERE NOT true|false` 85 | Please compare this version with the overridden test 86 | """ 87 | # testing "WHERE " renders a compatible expression 88 | boolean_table = self.tables.boolean_table 89 | 90 | with config.db.begin() as conn: 91 | conn.execute( 92 | boolean_table.insert(), 93 | [ 94 | {"id": 1, "value": True, "unconstrained_value": True}, 95 | {"id": 2, "value": False, "unconstrained_value": False}, 96 | ], 97 | ) 98 | 99 | eq_( 100 | conn.scalar( 101 | select(boolean_table.c.id).where(boolean_table.c.value) 102 | ), 103 | 1, 104 | ) 105 | eq_( 106 | conn.scalar( 107 | select(boolean_table.c.id).where( 108 | boolean_table.c.unconstrained_value 109 | ) 110 | ), 111 | 1, 112 | ) 113 | 114 | 115 | class CompoundSelectTest(_CompoundSelectTest): 116 | 117 | @testing.skip("databend") 118 | def test_limit_offset_aliased_selectable_in_unions(self): 119 | pass 120 | 121 | @testing.skip("databend") 122 | def test_limit_offset_selectable_in_unions(self): 123 | pass 124 | 125 | @testing.skip("databend") 126 | def test_limit_offset_in_unions_from_alias(self): 127 | pass 128 | 129 | 130 | class DeprecatedCompoundSelectTest(_DeprecatedCompoundSelectTest): 131 | @testing.skip("databend") 132 | def test_limit_offset_aliased_selectable_in_unions(self): 133 | pass 134 | 135 | @testing.skip("databend") 136 | def test_limit_offset_selectable_in_unions(self): 137 | pass 138 | 139 | @testing.skip("databend") 140 | def test_limit_offset_in_unions_from_alias(self): 141 | pass 142 | 143 | 144 | class HasIndexTest(_HasIndexTest): 145 | __requires__ = ('index_reflection',) 146 | 147 | 148 | class InsertBehaviorTest(_InsertBehaviorTest): 149 | @testing.skip("databend") # required autoinc columns 150 | def test_insert_from_select_autoinc(self, connection): 151 | pass 152 | 153 | @testing.skip("databend") # required autoinc columns 154 | def test_insert_from_select_autoinc_no_rows(self, connection): 155 | pass 156 | 157 | @testing.skip("databend") # required autoinc columns 158 | def test_no_results_for_non_returning_insert(self, connection): 159 | pass 160 | 161 | 162 | class LikeFunctionsTest(_LikeFunctionsTest): 163 | 164 | @testing.skip("databend") 165 | def test_contains_autoescape(self): 166 | pass 167 | 168 | @testing.skip("databend") 169 | def test_contains_escape(self): 170 | pass 171 | 172 | @testing.skip("databend") 173 | def test_contains_autoescape_escape(self): 174 | pass 175 | 176 | @testing.skip("databend") 177 | def test_endswith_autoescape(self): 178 | pass 179 | 180 | @testing.skip("databend") 181 | def test_endswith_escape(self): 182 | pass 183 | 184 | @testing.skip("databend") 185 | def test_endswith_autoescape_escape(self): 186 | pass 187 | 188 | @testing.skip("databend") 189 | def test_startswith_autoescape(self): 190 | pass 191 | 192 | @testing.skip("databend") 193 | def test_startswith_escape(self): 194 | pass 195 | 196 | @testing.skip("databend") 197 | def test_startswith_autoescape_escape(self): 198 | pass 199 | 200 | 201 | class LongNameBlowoutTest(_LongNameBlowoutTest): 202 | __requires__ = ("index_reflection",) # This will do to make it skip for now 203 | 204 | 205 | class QuotedNameArgumentTest(_QuotedNameArgumentTest): 206 | def quote_fixtures(fn): 207 | return testing.combinations( 208 | ("quote ' one",), 209 | ('quote " two', testing.requires.symbol_names_w_double_quote), 210 | )(fn) 211 | 212 | @quote_fixtures 213 | @testing.skip("databend") 214 | def test_get_pk_constraint(self, name): 215 | pass 216 | 217 | @quote_fixtures 218 | @testing.skip("databend") 219 | def test_get_foreign_keys(self, name): 220 | pass 221 | 222 | @quote_fixtures 223 | @testing.skip("databend") 224 | def test_get_indexes(self, name): 225 | pass 226 | 227 | 228 | class JoinTest(_JoinTest): 229 | __requires__ = ("foreign_keys",) 230 | 231 | if version.parse(sqlalchemy.__version__) >= version.parse('2.0.0') and version.parse(sqlalchemy.__version__) < version.parse('2.0.42'): 232 | class BizarroCharacterFKResolutionTest(_BizarroCharacterFKResolutionTest): 233 | __requires__ = ("foreign_keys",) 234 | 235 | 236 | class BinaryTest(_BinaryTest): 237 | 238 | # ToDo - get this working, failing because cannot substitute bytes parameter into sql statement 239 | # CREATE TABLE binary_test (x binary not null) 240 | # INSERT INTO binary_test (x) values (b'7\xe7\x9f') ??? 241 | # It is possible to do this 242 | # INSERT INTO binary_test (x) values (TO_BINARY('7\xe7\x9f')) 243 | # but that's not really a solution I don't think 244 | @testing.skip("databend") 245 | def test_binary_roundtrip(self): 246 | pass 247 | 248 | @testing.skip("databend") 249 | def test_pickle_roundtrip(self): 250 | pass 251 | 252 | 253 | class ServerSideCursorsTest(_ServerSideCursorsTest): 254 | 255 | def _is_server_side(self, cursor): 256 | # ToDo - requires implementation of `stream_results` option, so True always for now 257 | if self.engine.dialect.driver == "databend": 258 | return True 259 | return super() 260 | 261 | # ToDo - The commented out testing combinations here should be reviewed when `stream_results` is implemented 262 | @testing.combinations( 263 | ("global_string", True, "select 1", True), 264 | ("global_text", True, text("select 1"), True), 265 | ("global_expr", True, select(1), True), 266 | # ("global_off_explicit", False, text("select 1"), False), 267 | ( 268 | "stmt_option", 269 | False, 270 | select(1).execution_options(stream_results=True), 271 | True, 272 | ), 273 | # ( 274 | # "stmt_option_disabled", 275 | # True, 276 | # select(1).execution_options(stream_results=False), 277 | # False, 278 | # ), 279 | ("for_update_expr", True, select(1).with_for_update(), True), 280 | # TODO: need a real requirement for this, or dont use this test 281 | # ( 282 | # "for_update_string", 283 | # True, 284 | # "SELECT 1 FOR UPDATE", 285 | # True, 286 | # testing.skip_if(["sqlite", "mssql"]), 287 | # ), 288 | # ("text_no_ss", False, text("select 42"), False), 289 | ( 290 | "text_ss_option", 291 | False, 292 | text("select 42").execution_options(stream_results=True), 293 | True, 294 | ), 295 | id_="iaaa", 296 | argnames="engine_ss_arg, statement, cursor_ss_status", 297 | ) 298 | def test_ss_cursor_status( 299 | self, engine_ss_arg, statement, cursor_ss_status 300 | ): 301 | super() 302 | 303 | @testing.skip("databend") # ToDo - requires implementation of `stream_results` option 304 | def test_stmt_enabled_conn_option_disabled(self): 305 | pass 306 | 307 | @testing.skip("databend") # ToDo - requires implementation of `stream_results` option 308 | def test_aliases_and_ss(self): 309 | pass 310 | 311 | @testing.skip("databend") # Skipped because requires auto increment primary key 312 | def test_roundtrip_fetchall(self): 313 | pass 314 | 315 | @testing.skip("databend") # Skipped because requires auto increment primary key 316 | def test_roundtrip_fetchmany(self): 317 | pass 318 | 319 | if version.parse(sqlalchemy.__version__) >= version.parse('2.0.0'): 320 | class EnumTest(_EnumTest): 321 | __backend__ = True 322 | 323 | @testing.skip("databend") # Skipped because no supporting enums yet 324 | def test_round_trip_executemany(self, connection): 325 | pass 326 | 327 | 328 | class CTETest(_CTETest): 329 | @classmethod 330 | def define_tables(cls, metadata): 331 | Table( 332 | "some_table", 333 | metadata, 334 | Column("id", Integer, primary_key=True), 335 | Column("data", String(50)), 336 | Column("parent_id", Integer), # removed use of foreign key to get test to work 337 | ) 338 | 339 | Table( 340 | "some_other_table", 341 | metadata, 342 | Column("id", Integer, primary_key=True), 343 | Column("data", String(50)), 344 | Column("parent_id", Integer), 345 | ) 346 | 347 | 348 | class JSONTest(_JSONTest): 349 | @classmethod 350 | def define_tables(cls, metadata): 351 | Table( 352 | "data_table", 353 | metadata, 354 | Column("id", Integer), #, primary_key=True), # removed use of primary key to get test to work 355 | Column("name", String(30), nullable=False), 356 | Column("data", cls.datatype, nullable=False), 357 | Column("nulldata", cls.datatype(none_as_null=True)), 358 | ) 359 | 360 | # ToDo - this does not yet work 361 | def test_path_typed_comparison(self, datatype, value): 362 | pass 363 | 364 | 365 | class IntegerTest(_IntegerTest, fixtures.TablesTest): 366 | 367 | @classmethod 368 | def define_tables(cls, metadata): 369 | Table( 370 | "tiny_int_table", 371 | metadata, 372 | Column("id", TINYINT) 373 | ) 374 | 375 | def test_tinyint_write_and_read(self, connection): 376 | tiny_int_table = self.tables.tiny_int_table 377 | 378 | # Insert a value 379 | connection.execute( 380 | tiny_int_table.insert(), 381 | [{"id": 127}] # 127 is typically the maximum value for a signed TINYINT 382 | ) 383 | 384 | # Read the value back 385 | result = connection.execute(select(tiny_int_table.c.id)).scalar() 386 | 387 | # Verify the value 388 | eq_(result, 127) 389 | 390 | # Test with minimum value 391 | connection.execute( 392 | tiny_int_table.insert(), 393 | [{"id": -128}] # -128 is typically the minimum value for a signed TINYINT 394 | ) 395 | 396 | result = connection.execute(select(tiny_int_table.c.id).order_by(tiny_int_table.c.id)).first()[0] 397 | eq_(result, -128) 398 | 399 | def test_tinyint_overflow(self, connection): 400 | tiny_int_table = self.tables.tiny_int_table 401 | 402 | # This should raise an exception as it's outside the TINYINT range 403 | with assertions.expect_raises(Exception): # Replace with specific exception if known 404 | connection.execute( 405 | tiny_int_table.insert(), 406 | [{"id": 128}] # 128 is typically outside the range of a signed TINYINT 407 | ) 408 | 409 | with assertions.expect_raises(Exception): # Replace with specific exception if known 410 | connection.execute( 411 | tiny_int_table.insert(), 412 | [{"id": -129}] # -129 is typically outside the range of a signed TINYINT 413 | ) 414 | 415 | 416 | class BitmapTest(fixtures.TablesTest): 417 | 418 | @classmethod 419 | def define_tables(cls, metadata): 420 | Table( 421 | "bitmap_table", 422 | metadata, 423 | Column("id", Integer), 424 | Column("bitmap_data", BITMAP) 425 | ) 426 | 427 | """ 428 | Perform a simple test using Databend's bitmap data type to check 429 | that the bitmap data is correctly inserted and retrieved.' 430 | """ 431 | def test_bitmap_write_and_read(self, connection): 432 | bitmap_table = self.tables.bitmap_table 433 | 434 | # Insert a value 435 | connection.execute( 436 | bitmap_table.insert(), 437 | [{"id": 1, "bitmap_data": '1,2,3'}] 438 | ) 439 | 440 | # Read the value back 441 | result = connection.execute( 442 | select(bitmap_table.c.bitmap_data).where(bitmap_table.c.id == 1) 443 | ).scalar() 444 | 445 | # Verify the value 446 | eq_(result, ('1,2,3')) 447 | 448 | """ 449 | Perform a simple test using one of Databend's bitmap operations to check 450 | that the Bitmap data is correctly manipulated.' 451 | """ 452 | def test_bitmap_operations(self, connection): 453 | bitmap_table = self.tables.bitmap_table 454 | 455 | # Insert two values 456 | connection.execute( 457 | bitmap_table.insert(), 458 | [ 459 | {"id": 1, "bitmap_data": "1,4,5"}, 460 | {"id": 2, "bitmap_data": "4,5"} 461 | ] 462 | ) 463 | 464 | # Perform a bitmap AND operation and convert the result to a string 465 | result = connection.execute( 466 | select(func.to_string(func.bitmap_and( 467 | bitmap_table.c.bitmap_data, 468 | func.to_bitmap("3,4,5") 469 | ))).where(bitmap_table.c.id == 1) 470 | ).scalar() 471 | 472 | # Verify the result 473 | eq_(result, "4,5") 474 | 475 | 476 | class DoubleTest(fixtures.TablesTest): 477 | 478 | @classmethod 479 | def define_tables(cls, metadata): 480 | Table( 481 | "double_table", 482 | metadata, 483 | Column("id", Integer), 484 | Column("double_data", DOUBLE) 485 | ) 486 | 487 | def test_double_write_and_read(self, connection): 488 | double_table = self.tables.double_table 489 | 490 | # Insert a value 491 | connection.execute( 492 | double_table.insert(), 493 | [{"id": 1, "double_data": -1.7976931348623157E+308}] 494 | ) 495 | 496 | connection.execute( 497 | double_table.insert(), 498 | [{"id": 2, "double_data": 1.7976931348623157E+308}] 499 | ) 500 | 501 | # Read the value back 502 | result = connection.execute( 503 | select(double_table.c.double_data).where(double_table.c.id == 1) 504 | ).scalar() 505 | 506 | # Verify the value 507 | eq_(result, -1.7976931348623157E+308) 508 | 509 | # Read the value back 510 | result = connection.execute( 511 | select(double_table.c.double_data).where(double_table.c.id == 2) 512 | ).scalar() 513 | 514 | # Verify the value 515 | eq_(result, 1.7976931348623157E+308) 516 | 517 | 518 | def test_double_overflow(self, connection): 519 | double_table = self.tables.double_table 520 | 521 | # This should raise an exception as it's outside the DOUBLE range 522 | with assertions.expect_raises(Exception): 523 | connection.execute( 524 | double_table.insert(), 525 | [{"id": 3, "double_data": float('inf')}] 526 | ) 527 | 528 | with assertions.expect_raises(Exception): 529 | connection.execute( 530 | double_table.insert(), 531 | [{"id": 3, "double_data": float('-inf')}] 532 | ) 533 | 534 | 535 | class GeometryTest(fixtures.TablesTest): 536 | 537 | @classmethod 538 | def define_tables(cls, metadata): 539 | Table( 540 | "geometry_table", 541 | metadata, 542 | Column("id", Integer), 543 | Column("geometry_data", GEOMETRY) 544 | ) 545 | 546 | """ 547 | Perform a simple test using Databend's Geometry data type to check 548 | that the data is correctly inserted and retrieved.' 549 | """ 550 | def test_geometry_write_and_read(self, connection): 551 | geometry_table = self.tables.geometry_table 552 | 553 | # Insert a value 554 | connection.execute( 555 | geometry_table.insert(), 556 | [{"id": 1, "geometry_data": 'POINT(10 20)'}] 557 | ) 558 | connection.execute( 559 | geometry_table.insert(), 560 | [{"id": 2, "geometry_data": 'LINESTRING(10 20, 30 40, 50 60)'}] 561 | ) 562 | connection.execute( 563 | geometry_table.insert(), 564 | [{"id": 3, "geometry_data": 'POLYGON((10 20, 30 40, 50 60, 10 20))'}] 565 | ) 566 | connection.execute( 567 | geometry_table.insert(), 568 | [{"id": 4, "geometry_data": 'MULTIPOINT((10 20), (30 40), (50 60))'}] 569 | ) 570 | connection.execute( 571 | geometry_table.insert(), 572 | [{"id": 5, "geometry_data": 'MULTILINESTRING((10 20, 30 40), (50 60, 70 80))'}] 573 | ) 574 | connection.execute( 575 | geometry_table.insert(), 576 | [{"id": 6, "geometry_data": 'MULTIPOLYGON(((10 20, 30 40, 50 60, 10 20)), ((15 25, 25 35, 35 45, 15 25)))'}] 577 | ) 578 | connection.execute( 579 | geometry_table.insert(), 580 | [{"id": 7, "geometry_data": 'GEOMETRYCOLLECTION(POINT(10 20), LINESTRING(10 20, 30 40), POLYGON((10 20, 30 40, 50 60, 10 20)))'}] 581 | ) 582 | 583 | result = connection.execute( 584 | select(geometry_table.c.geometry_data).where(geometry_table.c.id == 1) 585 | ).scalar() 586 | eq_(result, ('{"type": "Point", "coordinates": [10,20]}')) 587 | result = connection.execute( 588 | select(geometry_table.c.geometry_data).where(geometry_table.c.id == 2) 589 | ).scalar() 590 | eq_(result, ('{"type": "LineString", "coordinates": [[10,20],[30,40],[50,60]]}')) 591 | result = connection.execute( 592 | select(geometry_table.c.geometry_data).where(geometry_table.c.id == 3) 593 | ).scalar() 594 | eq_(result, ('{"type": "Polygon", "coordinates": [[[10,20],[30,40],[50,60],[10,20]]]}')) 595 | result = connection.execute( 596 | select(geometry_table.c.geometry_data).where(geometry_table.c.id == 4) 597 | ).scalar() 598 | eq_(result, ('{"type": "MultiPoint", "coordinates": [[10,20],[30,40],[50,60]]}')) 599 | result = connection.execute( 600 | select(geometry_table.c.geometry_data).where(geometry_table.c.id == 5) 601 | ).scalar() 602 | eq_(result, ('{"type": "MultiLineString", "coordinates": [[[10,20],[30,40]],[[50,60],[70,80]]]}')) 603 | result = connection.execute( 604 | select(geometry_table.c.geometry_data).where(geometry_table.c.id == 6) 605 | ).scalar() 606 | eq_(result, ('{"type": "MultiPolygon", "coordinates": [[[[10,20],[30,40],[50,60],[10,20]]],[[[15,25],[25,35],[35,45],[15,25]]]]}')) 607 | result = connection.execute( 608 | select(geometry_table.c.geometry_data).where(geometry_table.c.id == 7) 609 | ).scalar() 610 | eq_(result, ('{"type": "GeometryCollection", "geometries": [{"type": "Point", "coordinates": [10,20]},{"type": "LineString", "coordinates": [[10,20],[30,40]]},{"type": "Polygon", "coordinates": [[[10,20],[30,40],[50,60],[10,20]]]}]}')) 611 | 612 | 613 | class GeographyTest(fixtures.TablesTest): 614 | 615 | @classmethod 616 | def define_tables(cls, metadata): 617 | Table( 618 | "geography_table", 619 | metadata, 620 | Column("id", Integer), 621 | Column("geography_data", GEOGRAPHY) 622 | ) 623 | 624 | """ 625 | Perform a simple test using Databend's Geography data type to check 626 | that the data is correctly inserted and retrieved.' 627 | """ 628 | def test_geography_write_and_read(self, connection): 629 | geography_table = self.tables.geography_table 630 | 631 | # Insert a value 632 | connection.execute( 633 | geography_table.insert(), 634 | [{"id": 1, "geography_data": 'POINT(10 20)'}] 635 | ) 636 | connection.execute( 637 | geography_table.insert(), 638 | [{"id": 2, "geography_data": 'LINESTRING(10 20, 30 40, 50 60)'}] 639 | ) 640 | connection.execute( 641 | geography_table.insert(), 642 | [{"id": 3, "geography_data": 'POLYGON((10 20, 30 40, 50 60, 10 20))'}] 643 | ) 644 | connection.execute( 645 | geography_table.insert(), 646 | [{"id": 4, "geography_data": 'MULTIPOINT((10 20), (30 40), (50 60))'}] 647 | ) 648 | connection.execute( 649 | geography_table.insert(), 650 | [{"id": 5, "geography_data": 'MULTILINESTRING((10 20, 30 40), (50 60, 70 80))'}] 651 | ) 652 | connection.execute( 653 | geography_table.insert(), 654 | [{"id": 6, "geography_data": 'MULTIPOLYGON(((10 20, 30 40, 50 60, 10 20)), ((15 25, 25 35, 35 45, 15 25)))'}] 655 | ) 656 | connection.execute( 657 | geography_table.insert(), 658 | [{"id": 7, "geography_data": 'GEOMETRYCOLLECTION(POINT(10 20), LINESTRING(10 20, 30 40), POLYGON((10 20, 30 40, 50 60, 10 20)))'}] 659 | ) 660 | 661 | result = connection.execute( 662 | select(geography_table.c.geography_data).where(geography_table.c.id == 1) 663 | ).scalar() 664 | eq_(result, ('{"type": "Point", "coordinates": [10,20]}')) 665 | result = connection.execute( 666 | select(geography_table.c.geography_data).where(geography_table.c.id == 2) 667 | ).scalar() 668 | eq_(result, ('{"type": "LineString", "coordinates": [[10,20],[30,40],[50,60]]}')) 669 | result = connection.execute( 670 | select(geography_table.c.geography_data).where(geography_table.c.id == 3) 671 | ).scalar() 672 | eq_(result, ('{"type": "Polygon", "coordinates": [[[10,20],[30,40],[50,60],[10,20]]]}')) 673 | result = connection.execute( 674 | select(geography_table.c.geography_data).where(geography_table.c.id == 4) 675 | ).scalar() 676 | eq_(result, ('{"type": "MultiPoint", "coordinates": [[10,20],[30,40],[50,60]]}')) 677 | result = connection.execute( 678 | select(geography_table.c.geography_data).where(geography_table.c.id == 5) 679 | ).scalar() 680 | eq_(result, ('{"type": "MultiLineString", "coordinates": [[[10,20],[30,40]],[[50,60],[70,80]]]}')) 681 | result = connection.execute( 682 | select(geography_table.c.geography_data).where(geography_table.c.id == 6) 683 | ).scalar() 684 | eq_(result, ('{"type": "MultiPolygon", "coordinates": [[[[10,20],[30,40],[50,60],[10,20]]],[[[15,25],[25,35],[35,45],[15,25]]]]}')) 685 | result = connection.execute( 686 | select(geography_table.c.geography_data).where(geography_table.c.id == 7) 687 | ).scalar() 688 | eq_(result, ('{"type": "GeometryCollection", "geometries": [{"type": "Point", "coordinates": [10,20]},{"type": "LineString", "coordinates": [[10,20],[30,40]]},{"type": "Polygon", "coordinates": [[[10,20],[30,40],[50,60],[10,20]]]}]}')) 689 | 690 | 691 | class HasSequenceTest(_HasSequenceTest): 692 | 693 | # ToDo - overridden other_seq definition due to lack of sequence ddl support for nominvalue nomaxvalue 694 | @classmethod 695 | def define_tables(cls, metadata): 696 | normalize_sequence(config, Sequence("user_id_seq", metadata=metadata)) 697 | normalize_sequence( 698 | config, 699 | Sequence( 700 | "other_seq", 701 | metadata=metadata, 702 | # nomaxvalue=True, 703 | # nominvalue=True, 704 | ), 705 | ) 706 | if testing.requires.schemas.enabled: 707 | #ToDo - omitted because Databend does not allow schema on sequence 708 | # normalize_sequence( 709 | # config, 710 | # Sequence( 711 | # "user_id_seq", schema=config.test_schema, metadata=metadata 712 | # ), 713 | # ) 714 | normalize_sequence( 715 | config, 716 | Sequence( 717 | "schema_seq", schema=config.test_schema, metadata=metadata 718 | ), 719 | ) 720 | Table( 721 | "user_id_table", 722 | metadata, 723 | Column("id", Integer, primary_key=True), 724 | ) 725 | 726 | @testing.skip("databend") # ToDo - requires definition of sequences with schema 727 | def test_has_sequence_remote_not_in_default(self, connection): 728 | eq_(inspect(connection).has_sequence("schema_seq"), False) 729 | 730 | @testing.skip("databend") # ToDo - requires definition of sequences with schema 731 | def test_get_sequence_names(self, connection): 732 | exp = {"other_seq", "user_id_seq"} 733 | 734 | res = set(inspect(connection).get_sequence_names()) 735 | is_true(res.intersection(exp) == exp) 736 | is_true("schema_seq" not in res) 737 | 738 | @testing.skip("databend") # ToDo - requires definition of sequences with schema 739 | @testing.requires.schemas 740 | def test_get_sequence_names_no_sequence_schema(self, connection): 741 | eq_( 742 | inspect(connection).get_sequence_names( 743 | schema=config.test_schema_2 744 | ), 745 | [], 746 | ) 747 | 748 | @testing.skip("databend") # ToDo - requires definition of sequences with schema 749 | @testing.requires.schemas 750 | def test_get_sequence_names_sequences_schema(self, connection): 751 | eq_( 752 | sorted( 753 | inspect(connection).get_sequence_names( 754 | schema=config.test_schema 755 | ) 756 | ), 757 | ["schema_seq", "user_id_seq"], 758 | ) 759 | --------------------------------------------------------------------------------