From f1cdc5246c7e893446fba293e10d5e81ec873527 Mon Sep 17 00:00:00 2001 From: "A. Fox" Date: Wed, 17 Jan 2024 06:48:29 -0800 Subject: [PATCH 01/12] add basic support for iceberg string and binary types --- sqlalchemy_redshift/dialect.py | 47 ++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/sqlalchemy_redshift/dialect.py b/sqlalchemy_redshift/dialect.py index ce755ea6..5d365443 100644 --- a/sqlalchemy_redshift/dialect.py +++ b/sqlalchemy_redshift/dialect.py @@ -290,6 +290,9 @@ class RedshiftImpl(postgresql.PostgresqlImpl): """ +MAX_VARCHAR_LENGTH = 65535 + + class RedshiftTypeEngine(TypeEngine): def _default_dialect(self, default=None): @@ -410,6 +413,47 @@ def get_dbapi_type(self, dbapi): return dbapi.HLLSKETCH +class IcebergString(sa.types.TypeDecorator): + impl = sa.types.String + + def load_dialect_impl(self, dialect): + return sa.dialects.postgresql.VARCHAR(length=MAX_VARCHAR_LENGTH) + + +class IcebergBinary(sa.types.TypeDecorator): + impl = sa.types.LargeBinary + + def process_bind_param(self, value, dialect): + if value is None: + return value + + encoding = getattr(dialect, 'encoding', 'utf-8') + + if isinstance(value, bytes): + return value + + if isinstance(value, str): + return value.encode(encoding) + + return value + + def result_processor(self, dialect, coltype): + def process(value): + if value is None: + return value + + if isinstance(value, bytes): + return value + + if isinstance(value, str): + encoding = getattr(dialect, 'encoding', 'utf-8') + return value.encode(encoding) + + raise TypeError(f"Unexpected type for value in result_processor.process: {type(value)}") + + return process + + # Mapping for database schema inspection of Amazon Redshift datatypes REDSHIFT_ISCHEMA_NAMES = { "geometry": GEOMETRY, @@ -417,6 +461,9 @@ def get_dbapi_type(self, dbapi): "time with time zone": TIMETZ, "timestamp with time zone": TIMESTAMPTZ, "hllsketch": HLLSKETCH, + # iceberg types + "string": IcebergString, + "binary": IcebergBinary, } From 39b95520cf9d7356c0b56188a70702bc2ea2a5d2 Mon Sep 17 00:00:00 2001 From: "A. Fox" Date: Wed, 17 Jan 2024 08:18:36 -0800 Subject: [PATCH 02/12] formatting and tests --- sqlalchemy_redshift/dialect.py | 7 +++++-- tests/test_column_loading.py | 31 ++++++++++++++++++++++++++++++- 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/sqlalchemy_redshift/dialect.py b/sqlalchemy_redshift/dialect.py index 5d365443..9ab74a24 100644 --- a/sqlalchemy_redshift/dialect.py +++ b/sqlalchemy_redshift/dialect.py @@ -422,7 +422,7 @@ def load_dialect_impl(self, dialect): class IcebergBinary(sa.types.TypeDecorator): impl = sa.types.LargeBinary - + def process_bind_param(self, value, dialect): if value is None: return value @@ -449,7 +449,10 @@ def process(value): encoding = getattr(dialect, 'encoding', 'utf-8') return value.encode(encoding) - raise TypeError(f"Unexpected type for value in result_processor.process: {type(value)}") + raise TypeError( + "Unexpected type for value in result_processor.process: ", + type(value) + ) return process diff --git a/tests/test_column_loading.py b/tests/test_column_loading.py index 349a31cc..cf5408a3 100644 --- a/tests/test_column_loading.py +++ b/tests/test_column_loading.py @@ -5,7 +5,8 @@ from sqlalchemy.types import NullType, VARCHAR from sqlalchemy_redshift.dialect import ( - RedshiftDialect_psycopg2, RedshiftDialect_psycopg2cffi + RedshiftDialect_psycopg2, RedshiftDialect_psycopg2cffi, + IcebergString, IcebergBinary ) sa_version = Version(sa.__version__) @@ -47,3 +48,31 @@ def test_varchar_as_nulltype(self): identity=None ) assert isinstance(varchar_info['type'], VARCHAR) + + iceberg_string_info = dialect._get_column_info( + name='Iceberg String Column', + format_type='string', + default=None, + notnull=False, + domains={}, + enums=[], + schema='default', + encode='', + comment='test column', + identity=None + ) + assert isinstance(iceberg_string_info['type'], IcebergString) + + iceberg_binary_info = dialect._get_column_info( + name='Iceberg Binary Column', + format_type='binary', + default=None, + notnull=False, + domains={}, + enums=[], + schema='default', + encode='', + comment='test column', + identity=None + ) + assert isinstance(iceberg_binary_info['type'], IcebergBinary) From 3839caf3198d63a9a0465a139bb5a86675fb3fed Mon Sep 17 00:00:00 2001 From: "A. Fox" Date: Wed, 17 Jan 2024 08:59:11 -0800 Subject: [PATCH 03/12] Column names should be uppercase since they are database-specific --- sqlalchemy_redshift/dialect.py | 8 ++++---- tests/test_column_loading.py | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/sqlalchemy_redshift/dialect.py b/sqlalchemy_redshift/dialect.py index 9ab74a24..b5c25e52 100644 --- a/sqlalchemy_redshift/dialect.py +++ b/sqlalchemy_redshift/dialect.py @@ -413,14 +413,14 @@ def get_dbapi_type(self, dbapi): return dbapi.HLLSKETCH -class IcebergString(sa.types.TypeDecorator): +class ICEBERG_STRING(sa.types.TypeDecorator): impl = sa.types.String def load_dialect_impl(self, dialect): return sa.dialects.postgresql.VARCHAR(length=MAX_VARCHAR_LENGTH) -class IcebergBinary(sa.types.TypeDecorator): +class ICEBERG_BINARY(sa.types.TypeDecorator): impl = sa.types.LargeBinary def process_bind_param(self, value, dialect): @@ -465,8 +465,8 @@ def process(value): "timestamp with time zone": TIMESTAMPTZ, "hllsketch": HLLSKETCH, # iceberg types - "string": IcebergString, - "binary": IcebergBinary, + "string": ICEBERG_STRING, + "binary": ICEBERG_BINARY, } diff --git a/tests/test_column_loading.py b/tests/test_column_loading.py index cf5408a3..ccdd35fb 100644 --- a/tests/test_column_loading.py +++ b/tests/test_column_loading.py @@ -6,7 +6,7 @@ from sqlalchemy_redshift.dialect import ( RedshiftDialect_psycopg2, RedshiftDialect_psycopg2cffi, - IcebergString, IcebergBinary + ICEBERG_STRING, ICEBERG_BINARY ) sa_version = Version(sa.__version__) @@ -61,7 +61,7 @@ def test_varchar_as_nulltype(self): comment='test column', identity=None ) - assert isinstance(iceberg_string_info['type'], IcebergString) + assert isinstance(iceberg_string_info['type'], ICEBERG_STRING) iceberg_binary_info = dialect._get_column_info( name='Iceberg Binary Column', @@ -75,4 +75,4 @@ def test_varchar_as_nulltype(self): comment='test column', identity=None ) - assert isinstance(iceberg_binary_info['type'], IcebergBinary) + assert isinstance(iceberg_binary_info['type'], ICEBERG_BINARY) From e70fa20acc96d6b8e37a13a6d1534ed6bdabbf42 Mon Sep 17 00:00:00 2001 From: "A. Fox" Date: Wed, 17 Jan 2024 09:02:50 -0800 Subject: [PATCH 04/12] consistency --- sqlalchemy_redshift/dialect.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/sqlalchemy_redshift/dialect.py b/sqlalchemy_redshift/dialect.py index b5c25e52..ef9f9e5e 100644 --- a/sqlalchemy_redshift/dialect.py +++ b/sqlalchemy_redshift/dialect.py @@ -427,12 +427,11 @@ def process_bind_param(self, value, dialect): if value is None: return value - encoding = getattr(dialect, 'encoding', 'utf-8') - if isinstance(value, bytes): return value if isinstance(value, str): + encoding = getattr(dialect, 'encoding', 'utf-8') return value.encode(encoding) return value From 26ffea3888ea279378395add67921dea7951986f Mon Sep 17 00:00:00 2001 From: "A. Fox" Date: Thu, 18 Jan 2024 12:54:42 -0800 Subject: [PATCH 05/12] add python_type --- sqlalchemy_redshift/dialect.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/sqlalchemy_redshift/dialect.py b/sqlalchemy_redshift/dialect.py index ef9f9e5e..adebf6e7 100644 --- a/sqlalchemy_redshift/dialect.py +++ b/sqlalchemy_redshift/dialect.py @@ -416,13 +416,20 @@ def get_dbapi_type(self, dbapi): class ICEBERG_STRING(sa.types.TypeDecorator): impl = sa.types.String + @property + def python_type(self): + return str + def load_dialect_impl(self, dialect): return sa.dialects.postgresql.VARCHAR(length=MAX_VARCHAR_LENGTH) - class ICEBERG_BINARY(sa.types.TypeDecorator): impl = sa.types.LargeBinary + @property + def python_type(self): + return bytes + def process_bind_param(self, value, dialect): if value is None: return value From 487173653d91c7c50acea19647873a697e0d319f Mon Sep 17 00:00:00 2001 From: "A. Fox" Date: Fri, 19 Jan 2024 06:53:07 -0800 Subject: [PATCH 06/12] Update CHANGES.rst --- CHANGES.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGES.rst b/CHANGES.rst index 3eea33d0..73aa19a0 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -1,7 +1,8 @@ 0.8.15 (unreleased) ------------------- -- Nothing changed yet. +- Add support for Iceberg string and binary types + (`Pull #297 `_) 0.8.14 (2023-04-07) From 7a57ce200692b3f0d6c8db4501ba5b4d3727e3ef Mon Sep 17 00:00:00 2001 From: "A. Fox" Date: Fri, 19 Jan 2024 07:05:41 -0800 Subject: [PATCH 07/12] Add docstrings --- sqlalchemy_redshift/dialect.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/sqlalchemy_redshift/dialect.py b/sqlalchemy_redshift/dialect.py index adebf6e7..7ebf4f7e 100644 --- a/sqlalchemy_redshift/dialect.py +++ b/sqlalchemy_redshift/dialect.py @@ -414,6 +414,13 @@ def get_dbapi_type(self, dbapi): class ICEBERG_STRING(sa.types.TypeDecorator): + """ + A custom SQLAlchemy type decorator for representing iceberg strings. + + This type decorator is used to represent iceberg strings in a + Redshift/PostgreSQL database using SQLAlchemy. + """ + impl = sa.types.String @property @@ -423,13 +430,22 @@ def python_type(self): def load_dialect_impl(self, dialect): return sa.dialects.postgresql.VARCHAR(length=MAX_VARCHAR_LENGTH) + class ICEBERG_BINARY(sa.types.TypeDecorator): + """ + A custom SQLAlchemy type decorator for storing/querying binary data in an + Iceberg database. + + This type decorator is used to represent iceberg binary data in a + Redshift/PostgreSQL database using SQLAlchemy. + """ + impl = sa.types.LargeBinary @property def python_type(self): return bytes - + def process_bind_param(self, value, dialect): if value is None: return value From 11a7d323062183576d0cd0c2912b8051de42e3c8 Mon Sep 17 00:00:00 2001 From: "A. Fox" Date: Fri, 19 Jan 2024 07:09:18 -0800 Subject: [PATCH 08/12] update docstrings --- sqlalchemy_redshift/dialect.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/sqlalchemy_redshift/dialect.py b/sqlalchemy_redshift/dialect.py index 7ebf4f7e..16243fb5 100644 --- a/sqlalchemy_redshift/dialect.py +++ b/sqlalchemy_redshift/dialect.py @@ -415,9 +415,7 @@ def get_dbapi_type(self, dbapi): class ICEBERG_STRING(sa.types.TypeDecorator): """ - A custom SQLAlchemy type decorator for representing iceberg strings. - - This type decorator is used to represent iceberg strings in a + ICEBERG_STRING is used to represent iceberg strings in a Redshift/PostgreSQL database using SQLAlchemy. """ @@ -433,10 +431,7 @@ def load_dialect_impl(self, dialect): class ICEBERG_BINARY(sa.types.TypeDecorator): """ - A custom SQLAlchemy type decorator for storing/querying binary data in an - Iceberg database. - - This type decorator is used to represent iceberg binary data in a + ICEBERG_BINARY is used to represent iceberg binary data in a Redshift/PostgreSQL database using SQLAlchemy. """ From 41e3d1d9a757179c315ca2f6fbd9aa15230b1008 Mon Sep 17 00:00:00 2001 From: "A. Fox" Date: Tue, 28 Jul 2026 13:23:34 -0700 Subject: [PATCH 09/12] feat: initial attempt. tests passing. --- .envrc | 1 + .gitignore | 2 + requirements-docs.txt | 7 +- setup.cfg | 13 +- sqlalchemy_redshift/__init__.py | 32 +- sqlalchemy_redshift/commands.py | 610 ++++++++++++++----------- sqlalchemy_redshift/dialect.py | 778 +++++++++++++++++++------------- tests/test_reflection_views.py | 38 +- tox.ini | 24 +- 9 files changed, 879 insertions(+), 626 deletions(-) create mode 100644 .envrc diff --git a/.envrc b/.envrc new file mode 100644 index 00000000..3b64139c --- /dev/null +++ b/.envrc @@ -0,0 +1 @@ +use pyenv 3.10.9 diff --git a/.gitignore b/.gitignore index efcc2a3d..9d610345 100644 --- a/.gitignore +++ b/.gitignore @@ -59,3 +59,5 @@ target/ # IDE .idea/ + +.direnv/ diff --git a/requirements-docs.txt b/requirements-docs.txt index da0a3e3f..cfded6c1 100644 --- a/requirements-docs.txt +++ b/requirements-docs.txt @@ -1,5 +1,4 @@ -e . -sphinx==1.6.3 -numpydoc==0.6.0 -psycopg2-binary==2.9.1 -jinja2<3.1.0 +sphinx==7.4.7 +numpydoc==1.8.0 +psycopg2-binary==2.9.11 diff --git a/setup.cfg b/setup.cfg index 3c6e79cf..fa270ee4 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,2 +1,13 @@ [bdist_wheel] -universal=1 +universal=0 + +[isort] +profile = black +line_length = 88 +known_first_party = sqlalchemy_redshift +force_sort_within_sections = true +split_on_trailing_comma = true + +[flake8] +max-line-length = 88 +extend-ignore = E203 diff --git a/sqlalchemy_redshift/__init__.py b/sqlalchemy_redshift/__init__.py index 934002cb..f3c350f2 100644 --- a/sqlalchemy_redshift/__init__.py +++ b/sqlalchemy_redshift/__init__.py @@ -1,31 +1,31 @@ -from pkg_resources import DistributionNotFound, get_distribution, parse_version +from importlib.metadata import PackageNotFoundError, version -for package in ['psycopg2', 'psycopg2-binary', 'psycopg2cffi']: +from packaging.version import Version + +for package in ["psycopg2", "psycopg2-binary", "psycopg2cffi"]: try: - if get_distribution(package).parsed_version < parse_version('2.5'): - raise ImportError('Minimum required version for psycopg2 is 2.5') + if Version(version(package)) < Version("2.5"): + raise ImportError("Minimum required version for psycopg2 is 2.5") break - except DistributionNotFound: + except PackageNotFoundError: pass -__version__ = get_distribution('sqlalchemy-redshift').version +__version__ = version("sqlalchemy-redshift") from sqlalchemy.dialects import registry # noqa +registry.register("redshift", "sqlalchemy_redshift.dialect", "RedshiftDialect_psycopg2") registry.register( - "redshift", "sqlalchemy_redshift.dialect", - "RedshiftDialect_psycopg2" -) -registry.register( - "redshift.psycopg2", "sqlalchemy_redshift.dialect", - "RedshiftDialect_psycopg2" + "redshift.psycopg2", "sqlalchemy_redshift.dialect", "RedshiftDialect_psycopg2" ) registry.register( - 'redshift+psycopg2cffi', 'sqlalchemy_redshift.dialect', - 'RedshiftDialect_psycopg2cffi', + "redshift+psycopg2cffi", + "sqlalchemy_redshift.dialect", + "RedshiftDialect_psycopg2cffi", ) registry.register( - "redshift+redshift_connector", "sqlalchemy_redshift.dialect", - "RedshiftDialect_redshift_connector" + "redshift+redshift_connector", + "sqlalchemy_redshift.dialect", + "RedshiftDialect_redshift_connector", ) diff --git a/sqlalchemy_redshift/commands.py b/sqlalchemy_redshift/commands.py index a3590d78..f0e6ecc5 100644 --- a/sqlalchemy_redshift/commands.py +++ b/sqlalchemy_redshift/commands.py @@ -2,6 +2,7 @@ import numbers import re import warnings + try: from collections.abc import Iterable except ImportError: @@ -12,7 +13,6 @@ from sqlalchemy.ext import compiler as sa_compiler from sqlalchemy.sql import expression as sa_expression - # At the time of this implementation, no specification for a session token was # found. After looking at a few session tokens they appear to be the same as # the aws_secret_access_key pattern, but much longer. An example token can be @@ -23,49 +23,55 @@ # The pattern of IAM role ARNs can be found here: # http://docs.aws.amazon.com/general/latest/gr/aws-arns-and-namespaces.html#arn-syntax-iam -ACCESS_KEY_ID_RE = re.compile('[A-Z0-9]{20}') -SECRET_ACCESS_KEY_RE = re.compile('[A-Za-z0-9/+=]{40}') -TOKEN_RE = re.compile('[A-Za-z0-9/+=]+') -AWS_PARTITIONS = frozenset({'aws', 'aws-cn', 'aws-us-gov'}) -AWS_ACCOUNT_ID_RE = re.compile('[0-9]{12}') -IAM_ROLE_NAME_RE = re.compile('[A-Za-z0-9+=,.@\-_]{1,64}') # noqa -IAM_ROLE_ARN_RE = re.compile('arn:(aws|aws-cn|aws-us-gov):iam::' - '[0-9]{12}:role/[A-Za-z0-9+=,.@\-_]{1,64}') # noqa - - -def _process_aws_credentials(access_key_id=None, secret_access_key=None, - session_token=None, aws_partition='aws', - aws_account_id=None, iam_role_name=None, - iam_role_arns=None): +ACCESS_KEY_ID_RE = re.compile("[A-Z0-9]{20}") +SECRET_ACCESS_KEY_RE = re.compile("[A-Za-z0-9/+=]{40}") +TOKEN_RE = re.compile("[A-Za-z0-9/+=]+") +AWS_PARTITIONS = frozenset({"aws", "aws-cn", "aws-us-gov"}) +AWS_ACCOUNT_ID_RE = re.compile("[0-9]{12}") +IAM_ROLE_NAME_RE = re.compile(r"[A-Za-z0-9+=,.@\-_]{1,64}") # noqa +IAM_ROLE_ARN_RE = re.compile( + r"arn:(aws|aws-cn|aws-us-gov):iam::[0-9]{12}:role/[A-Za-z0-9+=,.@\-_]{1,64}" # noqa +) + + +def _process_aws_credentials( + access_key_id=None, + secret_access_key=None, + session_token=None, + aws_partition="aws", + aws_account_id=None, + iam_role_name=None, + iam_role_arns=None, +): uses_iam_role = aws_account_id is not None and iam_role_name is not None uses_iam_roles = iam_role_arns is not None uses_key = access_key_id is not None and secret_access_key is not None if uses_iam_role + uses_iam_roles + uses_key > 1: raise TypeError( - 'Either access key based credentials or role based credentials ' - 'should be specified, but not both' + "Either access key based credentials or role based credentials " + "should be specified, but not both" ) credentials = None if aws_account_id is not None and iam_role_name is not None: if aws_partition not in AWS_PARTITIONS: - raise ValueError('invalid AWS partition') + raise ValueError("invalid AWS partition") if not AWS_ACCOUNT_ID_RE.match(aws_account_id): raise ValueError( - 'invalid AWS account ID; does not match {pattern}'.format( + "invalid AWS account ID; does not match {pattern}".format( pattern=AWS_ACCOUNT_ID_RE.pattern, ) ) elif not IAM_ROLE_NAME_RE.match(iam_role_name): raise ValueError( - 'invalid IAM role name; does not match {pattern}'.format( + "invalid IAM role name; does not match {pattern}".format( pattern=IAM_ROLE_NAME_RE.pattern, ) ) - credentials = 'aws_iam_role=arn:{0}:iam::{1}:role/{2}'.format( + credentials = "aws_iam_role=arn:{0}:iam::{1}:role/{2}".format( aws_partition, aws_account_id, iam_role_name, @@ -75,32 +81,32 @@ def _process_aws_credentials(access_key_id=None, secret_access_key=None, if isinstance(iam_role_arns, str): iam_role_arns = [iam_role_arns] if not isinstance(iam_role_arns, list): - raise ValueError('iam_role_arns must be a list') + raise ValueError("iam_role_arns must be a list") for arn in iam_role_arns: if not IAM_ROLE_ARN_RE.match(arn): raise ValueError( - 'invalid AWS account ID; does not match {pattern}'.format( + "invalid AWS account ID; does not match {pattern}".format( pattern=IAM_ROLE_ARN_RE.pattern, ) ) - credentials = 'aws_iam_role=' + ','.join(iam_role_arns) + credentials = "aws_iam_role=" + ",".join(iam_role_arns) if access_key_id is not None and secret_access_key is not None: if not ACCESS_KEY_ID_RE.match(access_key_id): raise ValueError( - 'invalid access_key_id; does not match {pattern}'.format( + "invalid access_key_id; does not match {pattern}".format( pattern=ACCESS_KEY_ID_RE.pattern, ) ) if not SECRET_ACCESS_KEY_RE.match(secret_access_key): raise ValueError( - 'invalid secret_access_key; does not match {pattern}'.format( + "invalid secret_access_key; does not match {pattern}".format( pattern=SECRET_ACCESS_KEY_RE.pattern, ) ) - credentials = 'aws_access_key_id={0};aws_secret_access_key={1}'.format( + credentials = "aws_access_key_id={0};aws_secret_access_key={1}".format( access_key_id, secret_access_key, ) @@ -108,27 +114,26 @@ def _process_aws_credentials(access_key_id=None, secret_access_key=None, if session_token is not None: if not TOKEN_RE.match(session_token): raise ValueError( - 'invalid session_token; does not match {pattern}'.format( + "invalid session_token; does not match {pattern}".format( pattern=TOKEN_RE.pattern, ) ) - credentials += ';token={0}'.format(session_token) + credentials += ";token={0}".format(session_token) if credentials is None: raise TypeError( - 'Either access key based credentials or role based credentials ' - 'should be specified' + "Either access key based credentials or role based credentials " + "should be specified" ) return credentials def _process_fixed_width(spec): - return ','.join(('{0}:{1:d}'.format(col, width) for col, width in spec)) + return ",".join(("{0}:{1:d}".format(col, width) for col, width in spec)) -class _ExecutableClause(sa_expression.Executable, - sa_expression.ClauseElement): +class _ExecutableClause(sa_expression.Executable, sa_expression.ClauseElement): pass @@ -159,10 +164,10 @@ class AlterTableAppendCommand(_ExecutableClause): fill those columns with the default column value or NULL. Mutually exclusive with `ignore_extra`. """ + def __init__(self, source, target, ignore_extra=False, fill_target=False): if ignore_extra and fill_target: - raise ValueError( - '"ignore_extra" cannot be used with "fill_target".') + raise ValueError('"ignore_extra" cannot be used with "fill_target".') self.source = source self.target = target @@ -176,18 +181,17 @@ def visit_alter_table_append_command(element, compiler, **kw): Returns the actual SQL query for the AlterTableAppendCommand class. """ if element.ignore_extra: - fill_option = 'IGNOREEXTRA' + fill_option = "IGNOREEXTRA" elif element.fill_target: - fill_option = 'FILLTARGET' + fill_option = "FILLTARGET" else: - fill_option = '' + fill_option = "" - query_text = \ - 'ALTER TABLE {target} APPEND FROM {source} {fill_option}'.format( - target=compiler.preparer.format_table(element.target), - source=compiler.preparer.format_table(element.source), - fill_option=fill_option, - ) + query_text = "ALTER TABLE {target} APPEND FROM {source} {fill_option}".format( + target=compiler.preparer.format_table(element.target), + source=compiler.preparer.format_table(element.source), + fill_option=fill_option, + ) return compiler.process(sa.text(query_text), **kw) @@ -269,24 +273,38 @@ class UnloadFromSelect(_ExecutableClause): Indicates the type of file to unload to. """ - def __init__(self, select, unload_location, access_key_id=None, - secret_access_key=None, session_token=None, - aws_partition='aws', aws_account_id=None, iam_role_name=None, - manifest=False, delimiter=None, fixed_width=None, - encrypted=False, gzip=False, add_quotes=False, null=None, - escape=False, allow_overwrite=False, parallel=True, - header=False, region=None, max_file_size=None, - format=None, iam_role_arns=None): + def __init__( + self, + select, + unload_location, + access_key_id=None, + secret_access_key=None, + session_token=None, + aws_partition="aws", + aws_account_id=None, + iam_role_name=None, + manifest=False, + delimiter=None, + fixed_width=None, + encrypted=False, + gzip=False, + add_quotes=False, + null=None, + escape=False, + allow_overwrite=False, + parallel=True, + header=False, + region=None, + max_file_size=None, + format=None, + iam_role_arns=None, + ): if delimiter is not None and len(delimiter) != 1: - raise ValueError( - '"delimiter" parameter must be a single character' - ) + raise ValueError('"delimiter" parameter must be a single character') if header and fixed_width is not None: - raise ValueError( - "'header' cannot be used with 'fixed_width'" - ) + raise ValueError("'header' cannot be used with 'fixed_width'") credentials = _process_aws_credentials( access_key_id=access_key_id, @@ -342,87 +360,96 @@ def visit_unload_from_select(element, compiler, **kw): el = element if el.format is None: - format_ = '' + format_ = "" elif el.format == Format.csv: - format_ = 'FORMAT AS {}'.format(el.format.value) + format_ = "FORMAT AS {}".format(el.format.value) if el.delimiter is not None or el.fixed_width is not None: - raise ValueError( - 'CSV format cannot be used with delimiter or fixed_width') + raise ValueError("CSV format cannot be used with delimiter or fixed_width") elif el.format == Format.parquet: - format_ = 'FORMAT AS {}'.format(el.format.value) - if any(( - el.delimiter, el.fixed_width, el.add_quotes, el.escape, el.null, - el.header, el.gzip - )): + format_ = "FORMAT AS {}".format(el.format.value) + if any( + ( + el.delimiter, + el.fixed_width, + el.add_quotes, + el.escape, + el.null, + el.header, + el.gzip, + ) + ): raise ValueError( "Parquet format can't be used with `delimiter`, `fixed_width`," - ' `add_quotes`, `escape`, `null`, `header`, or `gzip`.' + " `add_quotes`, `escape`, `null`, `header`, or `gzip`." ) else: - raise ValueError( - 'Only CSV and Parquet formats are currently supported.' - ) + raise ValueError("Only CSV and Parquet formats are currently supported.") qs = template.format( - manifest='MANIFEST' if el.manifest else '', - header='HEADER' if el.header else '', + manifest="MANIFEST" if el.manifest else "", + header="HEADER" if el.header else "", format=format_, - delimiter=( - 'DELIMITER AS :delimiter' if el.delimiter is not None else '' - ), - encrypted='ENCRYPTED' if el.encrypted else '', - fixed_width='FIXEDWIDTH AS :fixed_width' if el.fixed_width else '', - gzip='GZIP' if el.gzip else '', - add_quotes='ADDQUOTES' if el.add_quotes else '', - escape='ESCAPE' if el.escape else '', - null='NULL AS :null_as' if el.null is not None else '', - allow_overwrite='ALLOWOVERWRITE' if el.allow_overwrite else '', - parallel='PARALLEL OFF' if not el.parallel else '', - region='REGION :region' if el.region is not None else '', + delimiter=("DELIMITER AS :delimiter" if el.delimiter is not None else ""), + encrypted="ENCRYPTED" if el.encrypted else "", + fixed_width="FIXEDWIDTH AS :fixed_width" if el.fixed_width else "", + gzip="GZIP" if el.gzip else "", + add_quotes="ADDQUOTES" if el.add_quotes else "", + escape="ESCAPE" if el.escape else "", + null="NULL AS :null_as" if el.null is not None else "", + allow_overwrite="ALLOWOVERWRITE" if el.allow_overwrite else "", + parallel="PARALLEL OFF" if not el.parallel else "", + region="REGION :region" if el.region is not None else "", max_file_size=( - 'MAXFILESIZE :max_file_size MB' - if el.max_file_size is not None else '' + "MAXFILESIZE :max_file_size MB" if el.max_file_size is not None else "" ), ) query = sa.text(qs) if el.delimiter is not None: - query = query.bindparams(sa.bindparam( - 'delimiter', value=element.delimiter, type_=sa.String, - )) + query = query.bindparams( + sa.bindparam( + "delimiter", + value=element.delimiter, + type_=sa.String, + ) + ) if el.fixed_width: - query = query.bindparams(sa.bindparam( - 'fixed_width', - value=_process_fixed_width(el.fixed_width), - type_=sa.String, - )) + query = query.bindparams( + sa.bindparam( + "fixed_width", + value=_process_fixed_width(el.fixed_width), + type_=sa.String, + ) + ) if el.null is not None: - query = query.bindparams(sa.bindparam( - 'null_as', value=el.null, type_=sa.String - )) + query = query.bindparams( + sa.bindparam("null_as", value=el.null, type_=sa.String) + ) if el.region is not None: - query = query.bindparams(sa.bindparam( - 'region', value=el.region, type_=sa.String - )) + query = query.bindparams( + sa.bindparam("region", value=el.region, type_=sa.String) + ) if el.max_file_size is not None: max_file_size_mib = float(el.max_file_size) / 1024 / 1024 - query = query.bindparams(sa.bindparam( - 'max_file_size', value=max_file_size_mib, type_=sa.Float - )) + query = query.bindparams( + sa.bindparam("max_file_size", value=max_file_size_mib, type_=sa.Float) + ) return compiler.process( query.bindparams( - sa.bindparam('credentials', value=el.credentials, type_=sa.String), + sa.bindparam("credentials", value=el.credentials, type_=sa.String), sa.bindparam( - 'unload_location', value=el.unload_location, type_=sa.String, + "unload_location", + value=el.unload_location, + type_=sa.String, ), sa.bindparam( - 'select', + "select", value=compiler.process( el.select, literal_binds=True, @@ -430,30 +457,30 @@ def visit_unload_from_select(element, compiler, **kw): type_=sa.String, ), ), - **kw + **kw, ) class Format(enum.Enum): - csv = 'CSV' - json = 'JSON' - avro = 'AVRO' - orc = 'ORC' - parquet = 'PARQUET' - fixed_width = 'FIXEDWIDTH' + csv = "CSV" + json = "JSON" + avro = "AVRO" + orc = "ORC" + parquet = "PARQUET" + fixed_width = "FIXEDWIDTH" class Compression(enum.Enum): - gzip = 'GZIP' - lzop = 'LZOP' - bzip2 = 'BZIP2' + gzip = "GZIP" + lzop = "LZOP" + bzip2 = "BZIP2" class Encoding(enum.Enum): - utf8 = 'UTF8' - utf16 = 'UTF16' - utf16le = 'UTF16LE' - utf16be = 'UTF16BE' + utf8 = "UTF8" + utf16 = "UTF16" + utf16le = "UTF16LE" + utf16be = "UTF16BE" def _check_enum(Enum, val): @@ -462,7 +489,7 @@ def _check_enum(Enum, val): cleaned = Enum(val) if cleaned is not val: - tpl = '{val!r} should be, {cleaned!r}, an instance of {Enum!r}' + tpl = "{val!r} should be, {cleaned!r}, an instance of {Enum!r}" msg = tpl.format(val=val, cleaned=cleaned, Enum=Enum) warnings.warn(msg, DeprecationWarning) @@ -610,21 +637,48 @@ class CopyCommand(_ExecutableClause): cluster isn't in the same region as the S3 bucket. """ - def __init__(self, to, data_location, access_key_id=None, - secret_access_key=None, session_token=None, - aws_partition='aws', aws_account_id=None, iam_role_name=None, - format=None, quote=None, - path_file='auto', delimiter=None, fixed_width=None, - compression=None, accept_any_date=False, - accept_inv_chars=None, blanks_as_null=False, date_format=None, - empty_as_null=False, encoding=None, escape=False, - explicit_ids=False, fill_record=False, - ignore_blank_lines=False, ignore_header=None, - dangerous_null_delimiter=None, remove_quotes=False, - roundec=False, time_format=None, trim_blanks=False, - truncate_columns=False, comp_rows=None, comp_update=None, - max_error=None, no_load=False, stat_update=None, - manifest=False, region=None, iam_role_arns=None): + def __init__( + self, + to, + data_location, + access_key_id=None, + secret_access_key=None, + session_token=None, + aws_partition="aws", + aws_account_id=None, + iam_role_name=None, + format=None, + quote=None, + path_file="auto", + delimiter=None, + fixed_width=None, + compression=None, + accept_any_date=False, + accept_inv_chars=None, + blanks_as_null=False, + date_format=None, + empty_as_null=False, + encoding=None, + escape=False, + explicit_ids=False, + fill_record=False, + ignore_blank_lines=False, + ignore_header=None, + dangerous_null_delimiter=None, + remove_quotes=False, + roundec=False, + time_format=None, + trim_blanks=False, + truncate_columns=False, + comp_rows=None, + comp_update=None, + max_error=None, + no_load=False, + stat_update=None, + manifest=False, + region=None, + iam_role_arns=None, + ): credentials = _process_aws_credentials( access_key_id=access_key_id, @@ -637,14 +691,11 @@ def __init__(self, to, data_location, access_key_id=None, ) if delimiter is not None and len(delimiter) != 1: - raise ValueError('"delimiter" parameter must be a single ' - 'character') + raise ValueError('"delimiter" parameter must be a single character') if ignore_header is not None: if not isinstance(ignore_header, numbers.Integral): - raise TypeError( - '"ignore_header" parameter should be an integer' - ) + raise TypeError('"ignore_header" parameter should be an integer') table = None columns = [] @@ -652,10 +703,8 @@ def __init__(self, to, data_location, access_key_id=None, for column in to: if table is not None and table != column.table: raise ValueError( - 'All columns must come from the same table: ' - '%s comes from %s not %s' % ( - column, column.table, table - ), + "All columns must come from the same table: " + "%s comes from %s not %s" % (column, column.table, table), ) columns.append(column) table = column.table @@ -710,189 +759,211 @@ def visit_copy_command(element, compiler, **kw): parameters = [] bindparams = [ sa.bindparam( - 'data_location', + "data_location", value=element.data_location, type_=sa.String, ), sa.bindparam( - 'credentials', + "credentials", value=element.credentials, type_=sa.String, ), ] if element.format == Format.csv: - format_ = 'FORMAT AS CSV' + format_ = "FORMAT AS CSV" if element.quote is not None: - format_ += ' QUOTE AS :quote_character' - bindparams.append(sa.bindparam( - 'quote_character', - value=element.quote, - type_=sa.String, - )) + format_ += " QUOTE AS :quote_character" + bindparams.append( + sa.bindparam( + "quote_character", + value=element.quote, + type_=sa.String, + ) + ) elif element.format == Format.json: - format_ = 'FORMAT AS JSON AS :json_option' - bindparams.append(sa.bindparam( - 'json_option', - value=element.path_file, - type_=sa.String, - )) + format_ = "FORMAT AS JSON AS :json_option" + bindparams.append( + sa.bindparam( + "json_option", + value=element.path_file, + type_=sa.String, + ) + ) elif element.format == Format.avro: - format_ = 'FORMAT AS AVRO AS :avro_option' - bindparams.append(sa.bindparam( - 'avro_option', - value=element.path_file, - type_=sa.String, - )) + format_ = "FORMAT AS AVRO AS :avro_option" + bindparams.append( + sa.bindparam( + "avro_option", + value=element.path_file, + type_=sa.String, + ) + ) elif element.format == Format.orc: - format_ = 'FORMAT AS ORC' + format_ = "FORMAT AS ORC" elif element.format == Format.parquet: - format_ = 'FORMAT AS PARQUET' + format_ = "FORMAT AS PARQUET" elif element.format == Format.fixed_width and element.fixed_width is None: raise sa_exc.CompileError( - "'fixed_width' argument required for format 'FIXEDWIDTH'.") + "'fixed_width' argument required for format 'FIXEDWIDTH'." + ) else: - format_ = '' + format_ = "" if element.delimiter is not None: - parameters.append('DELIMITER AS :delimiter_char') - bindparams.append(sa.bindparam( - 'delimiter_char', - value=element.delimiter, - type_=sa.String, - )) + parameters.append("DELIMITER AS :delimiter_char") + bindparams.append( + sa.bindparam( + "delimiter_char", + value=element.delimiter, + type_=sa.String, + ) + ) if element.fixed_width is not None: - parameters.append('FIXEDWIDTH AS :fixedwidth_spec') - bindparams.append(sa.bindparam( - 'fixedwidth_spec', - value=_process_fixed_width(element.fixed_width), - type_=sa.String, - )) + parameters.append("FIXEDWIDTH AS :fixedwidth_spec") + bindparams.append( + sa.bindparam( + "fixedwidth_spec", + value=_process_fixed_width(element.fixed_width), + type_=sa.String, + ) + ) if element.compression is not None: parameters.append(Compression(element.compression).value) if element.manifest: - parameters.append('MANIFEST') + parameters.append("MANIFEST") if element.accept_any_date: - parameters.append('ACCEPTANYDATE') + parameters.append("ACCEPTANYDATE") if element.accept_inv_chars is not None: - parameters.append('ACCEPTINVCHARS AS :replacement_char') - bindparams.append(sa.bindparam( - 'replacement_char', - value=element.accept_inv_chars, - type_=sa.String - )) + parameters.append("ACCEPTINVCHARS AS :replacement_char") + bindparams.append( + sa.bindparam( + "replacement_char", value=element.accept_inv_chars, type_=sa.String + ) + ) if element.blanks_as_null: - parameters.append('BLANKSASNULL') + parameters.append("BLANKSASNULL") if element.date_format is not None: - parameters.append('DATEFORMAT AS :dateformat_string') - bindparams.append(sa.bindparam( - 'dateformat_string', - value=element.date_format, - type_=sa.String, - )) + parameters.append("DATEFORMAT AS :dateformat_string") + bindparams.append( + sa.bindparam( + "dateformat_string", + value=element.date_format, + type_=sa.String, + ) + ) if element.empty_as_null: - parameters.append('EMPTYASNULL') + parameters.append("EMPTYASNULL") if element.encoding is not None: - parameters.append('ENCODING AS ' + Encoding(element.encoding).value) + parameters.append("ENCODING AS " + Encoding(element.encoding).value) if element.escape: - parameters.append('ESCAPE') + parameters.append("ESCAPE") if element.explicit_ids: - parameters.append('EXPLICIT_IDS') + parameters.append("EXPLICIT_IDS") if element.fill_record: - parameters.append('FILLRECORD') + parameters.append("FILLRECORD") if element.ignore_blank_lines: - parameters.append('IGNOREBLANKLINES') + parameters.append("IGNOREBLANKLINES") if element.ignore_header is not None: - parameters.append('IGNOREHEADER AS :number_rows') - bindparams.append(sa.bindparam( - 'number_rows', - value=element.ignore_header, - type_=sa.Integer, - )) + parameters.append("IGNOREHEADER AS :number_rows") + bindparams.append( + sa.bindparam( + "number_rows", + value=element.ignore_header, + type_=sa.Integer, + ) + ) if element.dangerous_null_delimiter is not None: parameters.append("NULL AS '%s'" % element.dangerous_null_delimiter) if element.remove_quotes: - parameters.append('REMOVEQUOTES') + parameters.append("REMOVEQUOTES") if element.roundec: - parameters.append('ROUNDEC') + parameters.append("ROUNDEC") if element.time_format is not None: - parameters.append('TIMEFORMAT AS :timeformat_string') - bindparams.append(sa.bindparam( - 'timeformat_string', - value=element.time_format, - type_=sa.String, - )) + parameters.append("TIMEFORMAT AS :timeformat_string") + bindparams.append( + sa.bindparam( + "timeformat_string", + value=element.time_format, + type_=sa.String, + ) + ) if element.trim_blanks: - parameters.append('TRIMBLANKS') + parameters.append("TRIMBLANKS") if element.truncate_columns: - parameters.append('TRUNCATECOLUMNS') + parameters.append("TRUNCATECOLUMNS") if element.comp_rows: - parameters.append('COMPROWS :numrows') - bindparams.append(sa.bindparam( - 'numrows', - value=element.comp_rows, - type_=sa.Integer, - )) + parameters.append("COMPROWS :numrows") + bindparams.append( + sa.bindparam( + "numrows", + value=element.comp_rows, + type_=sa.Integer, + ) + ) if element.comp_update: - parameters.append('COMPUPDATE ON') + parameters.append("COMPUPDATE ON") elif element.comp_update is not None: - parameters.append('COMPUPDATE OFF') + parameters.append("COMPUPDATE OFF") if element.max_error is not None: - parameters.append('MAXERROR AS :error_count') - bindparams.append(sa.bindparam( - 'error_count', - value=element.max_error, - type_=sa.Integer, - )) + parameters.append("MAXERROR AS :error_count") + bindparams.append( + sa.bindparam( + "error_count", + value=element.max_error, + type_=sa.Integer, + ) + ) if element.no_load: - parameters.append('NOLOAD') + parameters.append("NOLOAD") if element.stat_update: - parameters.append('STATUPDATE ON') + parameters.append("STATUPDATE ON") elif element.stat_update is not None: - parameters.append('STATUPDATE OFF') + parameters.append("STATUPDATE OFF") if element.region is not None: - parameters.append('REGION :region') - bindparams.append(sa.bindparam( - 'region', - value=element.region, - type_=sa.String - )) + parameters.append("REGION :region") + bindparams.append(sa.bindparam("region", value=element.region, type_=sa.String)) - columns = ' (%s)' % ', '.join( - compiler.preparer.format_column(column) for column in element.columns - ) if element.columns else '' + columns = ( + " (%s)" + % ", ".join( + compiler.preparer.format_column(column) for column in element.columns + ) + if element.columns + else "" + ) qs = qs.format( table=compiler.preparer.format_table(element.table), columns=columns, format=format_, - parameters='\n'.join(parameters) + parameters="\n".join(parameters), ) return compiler.process(sa.text(qs).bindparams(*bindparams), **kw) @@ -941,10 +1012,20 @@ class CreateLibraryCommand(_ExecutableClause): The AWS region where the library's S3 bucket is located, if the Redshift cluster isn't in the same region as the S3 bucket. """ - def __init__(self, library_name, location, access_key_id=None, - secret_access_key=None, session_token=None, - aws_account_id=None, iam_role_name=None, replace=False, - region=None, iam_role_arns=None): + + def __init__( + self, + library_name, + location, + access_key_id=None, + secret_access_key=None, + session_token=None, + aws_account_id=None, + iam_role_name=None, + replace=False, + region=None, + iam_role_arns=None, + ): self.library_name = library_name self.location = location self.credentials = _process_aws_credentials( @@ -973,28 +1054,32 @@ def visit_create_library_command(element, compiler, **kw): """ bindparams = [ sa.bindparam( - 'location', + "location", value=element.location, type_=sa.String, ), sa.bindparam( - 'credentials', + "credentials", value=element.credentials, type_=sa.String, ), ] if element.region is not None: - bindparams.append(sa.bindparam( - 'region', - value=element.region, - type_=sa.String, - )) + bindparams.append( + sa.bindparam( + "region", + value=element.region, + type_=sa.String, + ) + ) quoted_lib_name = compiler.preparer.quote_identifier(element.library_name) - query = query.format(name=quoted_lib_name, - or_replace='OR REPLACE' if element.replace else '', - region='REGION :region' if element.region else '') + query = query.format( + name=quoted_lib_name, + or_replace="OR REPLACE" if element.replace else "", + region="REGION :region" if element.region else "", + ) return compiler.process(sa.text(query).bindparams(*bindparams), **kw) @@ -1019,6 +1104,7 @@ class RefreshMaterializedView(_ExecutableClause): This can be included in any execute() statement. """ + def __init__(self, name): """ Builds the Executable/ClauseElement that represents the refresh command diff --git a/sqlalchemy_redshift/dialect.py b/sqlalchemy_redshift/dialect.py index 16243fb5..bba4f227 100644 --- a/sqlalchemy_redshift/dialect.py +++ b/sqlalchemy_redshift/dialect.py @@ -1,34 +1,54 @@ +from collections import defaultdict, namedtuple import importlib +from importlib.resources import files import json -import re -from collections import defaultdict, namedtuple from logging import getLogger +import re -import pkg_resources -import sqlalchemy as sa from packaging.version import Version +import sqlalchemy as sa from sqlalchemy import inspect from sqlalchemy.dialects.postgresql import DOUBLE_PRECISION -from sqlalchemy.dialects.postgresql.base import (PGCompiler, PGDDLCompiler, - PGDialect, PGExecutionContext, - PGIdentifierPreparer, - PGTypeCompiler) +from sqlalchemy.dialects.postgresql.base import ( + PGCompiler, + PGDDLCompiler, + PGDialect, + PGExecutionContext, + PGIdentifierPreparer, + PGTypeCompiler, +) from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2 from sqlalchemy.dialects.postgresql.psycopg2cffi import PGDialect_psycopg2cffi from sqlalchemy.engine import reflection from sqlalchemy.engine.default import DefaultDialect from sqlalchemy.ext.compiler import compiles -from sqlalchemy.sql.expression import (BinaryExpression, BooleanClauseList, - Delete) +from sqlalchemy.sql.expression import BinaryExpression, BooleanClauseList, Delete from sqlalchemy.sql.type_api import TypeEngine -from sqlalchemy.types import (BIGINT, BOOLEAN, CHAR, DATE, DECIMAL, INTEGER, - REAL, SMALLINT, TIMESTAMP, VARCHAR, NullType) +from sqlalchemy.types import ( + BIGINT, + BOOLEAN, + CHAR, + DATE, + DECIMAL, + INTEGER, + REAL, + SMALLINT, + TIMESTAMP, + VARCHAR, + NullType, +) -from .commands import (AlterTableAppendCommand, Compression, CopyCommand, - CreateLibraryCommand, Encoding, Format, - RefreshMaterializedView, UnloadFromSelect) -from .ddl import (CreateMaterializedView, DropMaterializedView, - get_table_attributes) +from .commands import ( + AlterTableAppendCommand, + Compression, + CopyCommand, + CreateLibraryCommand, + Encoding, + Format, + RefreshMaterializedView, + UnloadFromSelect, +) +from .ddl import CreateMaterializedView, DropMaterializedView, get_table_attributes sa_version = Version(sa.__version__) logger = getLogger(__name__) @@ -40,50 +60,59 @@ else: from alembic.ddl import postgresql from alembic.ddl.base import RenameTable - compiles(RenameTable, 'redshift')(postgresql.visit_rename_table) - if Version(alembic.__version__) >= Version('1.0.6'): + compiles(RenameTable, "redshift")(postgresql.visit_rename_table) + + if Version(alembic.__version__) >= Version("1.0.6"): from alembic.ddl.base import ColumnComment - compiles(ColumnComment, 'redshift')(postgresql.visit_column_comment) + + compiles(ColumnComment, "redshift")(postgresql.visit_column_comment) class RedshiftImpl(postgresql.PostgresqlImpl): - __dialect__ = 'redshift' + __dialect__ = "redshift" + # "Each dialect provides the full set of typenames supported by that backend # with its __all__ collection # https://docs.sqlalchemy.org/en/13/core/type_basics.html#vendor-specific-types __all__ = ( - 'SMALLINT', - 'INTEGER', - 'BIGINT', - 'DECIMAL', - 'REAL', - 'BOOLEAN', - 'CHAR', - 'DATE', - 'TIMESTAMP', - 'VARCHAR', - 'DOUBLE_PRECISION', - 'GEOMETRY', - 'SUPER', - 'TIMESTAMPTZ', - 'TIMETZ', - 'HLLSKETCH', - - 'RedshiftDialect', 'RedshiftDialect_psycopg2', - 'RedshiftDialect_psycopg2cffi', 'RedshiftDialect_redshift_connector', - - 'CopyCommand', 'UnloadFromSelect', 'Compression', - 'Encoding', 'Format', 'CreateLibraryCommand', 'AlterTableAppendCommand', - 'RefreshMaterializedView', - - 'CreateMaterializedView', 'DropMaterializedView' + "SMALLINT", + "INTEGER", + "BIGINT", + "DECIMAL", + "REAL", + "BOOLEAN", + "CHAR", + "DATE", + "TIMESTAMP", + "VARCHAR", + "DOUBLE_PRECISION", + "GEOMETRY", + "SUPER", + "TIMESTAMPTZ", + "TIMETZ", + "HLLSKETCH", + "RedshiftDialect", + "RedshiftDialect_psycopg2", + "RedshiftDialect_psycopg2cffi", + "RedshiftDialect_redshift_connector", + "CopyCommand", + "UnloadFromSelect", + "Compression", + "Encoding", + "Format", + "CreateLibraryCommand", + "AlterTableAppendCommand", + "RefreshMaterializedView", + "CreateMaterializedView", + "DropMaterializedView", ) # Regex for parsing and identity constraint out of adsrc, e.g.: # "identity"(445178, 0, '1,1'::text) -IDENTITY_RE = re.compile(r""" +IDENTITY_RE = re.compile( + r""" "identity" \( (?P-?\d+) ,\s @@ -92,20 +121,26 @@ class RedshiftImpl(postgresql.PostgresqlImpl): '(?P-?\d+),(?P-?\d+)' .* \) -""", re.VERBOSE) +""", + re.VERBOSE, +) # Regex for SQL identifiers (valid table and column names) -SQL_IDENTIFIER_RE = re.compile(r""" +SQL_IDENTIFIER_RE = re.compile( + r""" [_a-zA-Z][\w$]* # SQL standard identifier | # or (?:"[^"]+")+ # SQL delimited (quoted) identifier -""", re.VERBOSE) +""", + re.VERBOSE, +) # Regex for foreign key constraints, e.g.: # FOREIGN KEY(col1) REFERENCES othertable (col2) # See https://docs.aws.amazon.com/redshift/latest/dg/r_names.html # for a definition of valid SQL identifiers. -FOREIGN_KEY_RE = re.compile(r""" +FOREIGN_KEY_RE = re.compile( + r""" ^FOREIGN\ KEY \s* \( # FOREIGN KEY, arbitrary whitespace, literal '(' (?P # Start a group to capture the referring columns (?: # Start a non-capturing group @@ -129,11 +164,14 @@ class RedshiftImpl(postgresql.PostgresqlImpl): )+ # Close the non-capturing group; require at least one ) # Close the 'columns' group \s* \) # Arbitrary whitespace and literal ')' -""", re.VERBOSE) +""", + re.VERBOSE, +) # Regex for primary key constraints, e.g.: # PRIMARY KEY (col1, col2) -PRIMARY_KEY_RE = re.compile(r""" +PRIMARY_KEY_RE = re.compile( + r""" ^PRIMARY \s* KEY \s* \( # FOREIGN KEY, arbitrary whitespace, literal '(' (?P # Start a group to capture column names (?: @@ -145,38 +183,173 @@ class RedshiftImpl(postgresql.PostgresqlImpl): )+ # Close the non-capturing group; require at least one ) \s* \) \s* # Arbitrary whitespace and literal ')' -""", re.VERBOSE) +""", + re.VERBOSE, +) # Reserved words as extracted from Redshift docs. # See pull_reserved_words.sh at the top level of this repository # for the code used to generate this set. -RESERVED_WORDS = set([ - "aes128", "aes256", "all", "allowoverwrite", "analyse", "analyze", - "and", "any", "array", "as", "asc", "authorization", "az64", - "backup", "between", "binary", "blanksasnull", "both", "bytedict", - "bzip2", "case", "cast", "check", "collate", "column", "constraint", - "create", "credentials", "cross", "current_date", "current_time", - "current_timestamp", "current_user", "current_user_id", "default", - "deferrable", "deflate", "defrag", "delta", "delta32k", "desc", - "disable", "distinct", "do", "else", "emptyasnull", "enable", - "encode", "encrypt", "encryption", "end", "except", "explicit", - "false", "for", "foreign", "freeze", "from", "full", "globaldict256", - "globaldict64k", "grant", "group", "gzip", "having", "identity", - "ignore", "ilike", "in", "initially", "inner", "intersect", "into", - "is", "isnull", "join", "language", "leading", "left", "like", - "limit", "localtime", "localtimestamp", "lun", "luns", "lzo", "lzop", - "minus", "mostly16", "mostly32", "mostly8", "natural", "new", "not", - "notnull", "null", "nulls", "off", "offline", "offset", "oid", "old", - "on", "only", "open", "or", "order", "outer", "overlaps", "parallel", - "partition", "percent", "permissions", "pivot", "placing", "primary", - "raw", "readratio", "recover", "references", "respect", "rejectlog", - "resort", "restore", "right", "select", "session_user", "similar", - "snapshot", "some", "sysdate", "system", "table", "tag", "tdes", - "text255", "text32k", "then", "timestamp", "to", "top", "trailing", - "true", "truncatecolumns", "union", "unique", "unnest", "unpivot", - "user", "using", "verbose", "wallet", "when", "where", "with", - "without", -]) +RESERVED_WORDS = set( + [ + "aes128", + "aes256", + "all", + "allowoverwrite", + "analyse", + "analyze", + "and", + "any", + "array", + "as", + "asc", + "authorization", + "az64", + "backup", + "between", + "binary", + "blanksasnull", + "both", + "bytedict", + "bzip2", + "case", + "cast", + "check", + "collate", + "column", + "constraint", + "create", + "credentials", + "cross", + "current_date", + "current_time", + "current_timestamp", + "current_user", + "current_user_id", + "default", + "deferrable", + "deflate", + "defrag", + "delta", + "delta32k", + "desc", + "disable", + "distinct", + "do", + "else", + "emptyasnull", + "enable", + "encode", + "encrypt", + "encryption", + "end", + "except", + "explicit", + "false", + "for", + "foreign", + "freeze", + "from", + "full", + "globaldict256", + "globaldict64k", + "grant", + "group", + "gzip", + "having", + "identity", + "ignore", + "ilike", + "in", + "initially", + "inner", + "intersect", + "into", + "is", + "isnull", + "join", + "language", + "leading", + "left", + "like", + "limit", + "localtime", + "localtimestamp", + "lun", + "luns", + "lzo", + "lzop", + "minus", + "mostly16", + "mostly32", + "mostly8", + "natural", + "new", + "not", + "notnull", + "null", + "nulls", + "off", + "offline", + "offset", + "oid", + "old", + "on", + "only", + "open", + "or", + "order", + "outer", + "overlaps", + "parallel", + "partition", + "percent", + "permissions", + "pivot", + "placing", + "primary", + "raw", + "readratio", + "recover", + "references", + "respect", + "rejectlog", + "resort", + "restore", + "right", + "select", + "session_user", + "similar", + "snapshot", + "some", + "sysdate", + "system", + "table", + "tag", + "tdes", + "text255", + "text32k", + "then", + "timestamp", + "to", + "top", + "trailing", + "true", + "truncatecolumns", + "union", + "unique", + "unnest", + "unpivot", + "user", + "using", + "verbose", + "wallet", + "when", + "where", + "with", + "without", + ] +) REFLECTION_SQL = """\ SELECT @@ -294,7 +467,6 @@ class RedshiftImpl(postgresql.PostgresqlImpl): class RedshiftTypeEngine(TypeEngine): - def _default_dialect(self, default=None): """ Returns the default dialect used for TypeEngine compilation yielding @@ -317,7 +489,7 @@ class TIMESTAMPTZ(RedshiftTypeEngine, sa.dialects.postgresql.TIMESTAMP): https://docs.sqlalchemy.org/en/13/core/type_basics.html#vendor-specific-types """ - __visit_name__ = 'TIMESTAMPTZ' + __visit_name__ = "TIMESTAMPTZ" def __init__(self, timezone=True, precision=None): # timezone param must be present as it's provided in base class so the @@ -338,7 +510,7 @@ class TIMETZ(RedshiftTypeEngine, sa.dialects.postgresql.TIME): https://docs.sqlalchemy.org/en/13/core/type_basics.html#vendor-specific-types """ - __visit_name__ = 'TIMETZ' + __visit_name__ = "TIMETZ" def __init__(self, timezone=True, precision=None): # timezone param must be present as it's provided in base class so the @@ -357,7 +529,8 @@ class GEOMETRY(RedshiftTypeEngine, sa.dialects.postgresql.TEXT): https://docs.sqlalchemy.org/en/13/core/type_basics.html#vendor-specific-types """ - __visit_name__ = 'GEOMETRY' + + __visit_name__ = "GEOMETRY" def __init__(self): super(GEOMETRY, self).__init__() @@ -377,7 +550,7 @@ class SUPER(RedshiftTypeEngine, sa.dialects.postgresql.TEXT): https://docs.sqlalchemy.org/en/13/core/type_basics.html#vendor-specific-types """ - __visit_name__ = 'SUPER' + __visit_name__ = "SUPER" def __init__(self): super(SUPER, self).__init__() @@ -404,7 +577,8 @@ class HLLSKETCH(RedshiftTypeEngine, sa.dialects.postgresql.TEXT): https://docs.sqlalchemy.org/en/13/core/type_basics.html#vendor-specific-types """ - __visit_name__ = 'HLLSKETCH' + + __visit_name__ = "HLLSKETCH" def __init__(self): super(HLLSKETCH, self).__init__() @@ -449,7 +623,7 @@ def process_bind_param(self, value, dialect): return value if isinstance(value, str): - encoding = getattr(dialect, 'encoding', 'utf-8') + encoding = getattr(dialect, "encoding", "utf-8") return value.encode(encoding) return value @@ -463,12 +637,11 @@ def process(value): return value if isinstance(value, str): - encoding = getattr(dialect, 'encoding', 'utf-8') + encoding = getattr(dialect, "encoding", "utf-8") return value.encode(encoding) raise TypeError( - "Unexpected type for value in result_processor.process: ", - type(value) + "Unexpected type for value in result_processor.process: ", type(value) ) return process @@ -487,10 +660,11 @@ def process(value): } -class RelationKey(namedtuple('RelationKey', ('name', 'schema'))): +class RelationKey(namedtuple("RelationKey", ("name", "schema"))): """ Structured tuple of table/view name and schema name. """ + __slots__ = () def __new__(cls, name, schema=None, connection=None): @@ -511,10 +685,7 @@ def __str__(self): @staticmethod def _unquote(part): - if ( - part is not None and part.startswith('"') and - part.endswith('"') - ): + if part is not None and part.startswith('"') and part.endswith('"'): return part[1:-1] return part @@ -527,13 +698,11 @@ def unquoted(self): In particular, this happens for tables named as a keyword. """ return RelationKey( - RelationKey._unquote(self.name), - RelationKey._unquote(self.schema) + RelationKey._unquote(self.name), RelationKey._unquote(self.schema) ) class RedshiftCompiler(PGCompiler): - def visit_now_func(self, fn, **kw): return "SYSDATE" @@ -649,7 +818,7 @@ class RedshiftDDLCompiler(PGDDLCompiler): def post_create_table(self, table): kwargs = ["diststyle", "distkey", "sortkey", "interleaved_sortkey"] - info = table.dialect_options['redshift'] + info = table.dialect_options["redshift"] info = {key: info.get(key) for key in kwargs} return get_table_attributes(self.preparer, **info) @@ -675,33 +844,32 @@ def get_column_specification(self, column, **kwargs): def _fetch_redshift_column_attributes(self, column): text = "" - if sa_version >= Version('1.3.0'): - info = column.dialect_options['redshift'] + if sa_version >= Version("1.3.0"): + info = column.dialect_options["redshift"] else: - if not hasattr(column, 'info'): + if not hasattr(column, "info"): return text info = column.info - identity = info.get('identity') + identity = info.get("identity") if identity: text += " IDENTITY({0},{1})".format(identity[0], identity[1]) - encode = info.get('encode') + encode = info.get("encode") if encode: text += " ENCODE " + encode - distkey = info.get('distkey') + distkey = info.get("distkey") if distkey: text += " DISTKEY" - sortkey = info.get('sortkey') + sortkey = info.get("sortkey") if sortkey: text += " SORTKEY" return text class RedshiftTypeCompiler(PGTypeCompiler): - def visit_GEOMETRY(self, type_, **kw): return "GEOMETRY" @@ -731,7 +899,7 @@ class RedshiftDialectMixin(DefaultDialect): :class:`~sqlalchemy.engine.Inspector`. """ - name = 'redshift' + name = "redshift" max_identifier_length = 127 statement_compiler = RedshiftCompiler @@ -739,24 +907,26 @@ class RedshiftDialectMixin(DefaultDialect): preparer = RedshiftIdentifierPreparer type_compiler = RedshiftTypeCompiler construct_arguments = [ - (sa.schema.Index, { - "using": False, - "where": None, - "ops": {} - }), - (sa.schema.Table, { - "ignore_search_path": False, - "diststyle": None, - "distkey": None, - "sortkey": None, - "interleaved_sortkey": None, - }), - (sa.schema.Column, { - "encode": None, - "distkey": None, - "sortkey": None, - "identity": None, - }), + (sa.schema.Index, {"using": False, "where": None, "ops": {}}), + ( + sa.schema.Table, + { + "ignore_search_path": False, + "diststyle": None, + "distkey": None, + "sortkey": None, + "interleaved_sortkey": None, + }, + ), + ( + sa.schema.Column, + { + "encode": None, + "distkey": None, + "sortkey": None, + "identity": None, + }, + ), ] def __init__(self, *args, **kw): @@ -775,7 +945,7 @@ def ischema_names(self): """ return { **super(RedshiftDialectMixin, self).ischema_names, - **REDSHIFT_ISCHEMA_NAMES + **REDSHIFT_ISCHEMA_NAMES, } @reflection.cache @@ -793,10 +963,16 @@ def get_columns(self, connection, table_name, schema=None, **kw): columns = [] for col in cols: column_info = self._get_column_info( - name=col.name, format_type=col.format_type, - default=col.default, notnull=col.notnull, domains=domains, - enums=[], schema=col.schema, encode=col.encode, - comment=col.comment) + name=col.name, + format_type=col.format_type, + default=col.default, + notnull=col.notnull, + domains=domains, + enums=[], + schema=col.schema, + encode=col.encode, + comment=col.comment, + ) columns.append(column_info) return columns @@ -805,11 +981,10 @@ def has_table(self, connection, table_name, schema=None, **kw): if not schema: schema = inspect(connection).default_schema_name - info_cache = kw.get('info_cache') - table = self._get_all_relation_info(connection, - schema=schema, - table_name=table_name, - info_cache=info_cache) + info_cache = kw.get("info_cache") + table = self._get_all_relation_info( + connection, schema=schema, table_name=table_name, info_cache=info_cache + ) return True if table else False @@ -818,9 +993,11 @@ def get_check_constraints(self, connection, table_name, schema=None, **kw): table_oid = self.get_table_oid( connection, table_name, schema, info_cache=kw.get("info_cache") ) - table_oid = 'NULL' if not table_oid else table_oid + table_oid = "NULL" if not table_oid else table_oid - result = connection.execute(sa.text(""" + result = connection.execute( + sa.text( + """ SELECT cons.conname as name, pg_get_constraintdef(cons.oid) as src @@ -829,7 +1006,9 @@ def get_check_constraints(self, connection, table_name, schema=None, **kw): WHERE cons.conrelid = {} AND cons.contype = 'c' - """.format(table_oid))) + """.format(table_oid) + ) + ) ret = [] for name, src in result: # samples: @@ -839,16 +1018,14 @@ def get_check_constraints(self, connection, table_name, schema=None, **kw): # "CHECK (some_boolean_function(a))" # "CHECK (((a\n < 1)\n OR\n (a\n >= 5))\n)" - m = re.match( - r"^CHECK *\((.+)\)( NOT VALID)?$", src, flags=re.DOTALL - ) + m = re.match(r"^CHECK *\((.+)\)( NOT VALID)?$", src, flags=re.DOTALL) if not m: logger.warning(f"Could not parse CHECK constraint text: {src}") sqltext = "" else: - sqltext = re.compile( - r"^[\s\n]*\((.+)\)[\s\n]*$", flags=re.DOTALL - ).sub(r"\1", m.group(1)) + sqltext = re.compile(r"^[\s\n]*\((.+)\)[\s\n]*$", flags=re.DOTALL).sub( + r"\1", m.group(1) + ) entry = {"name": name, "sqltext": sqltext} if m and m.group(2): entry["dialect_options"] = {"not_valid": True} @@ -866,10 +1043,7 @@ def get_table_oid(self, connection, table_name, schema=None, **kw): sa.text( """ select '{schema_field}"{table_name}"'::regclass::oid; - """.format( - schema_field=schema_field, - table_name=table_name - ) + """.format(schema_field=schema_field, table_name=table_name) ) ) @@ -883,18 +1057,19 @@ def get_pk_constraint(self, connection, table_name, schema=None, **kw): Overrides interface :meth:`~sqlalchemy.engine.interfaces.Dialect.get_pk_constraint`. """ - constraints = self._get_redshift_constraints(connection, table_name, - schema, **kw) - pk_constraints = [c for c in constraints if c.contype == 'p'] + constraints = self._get_redshift_constraints( + connection, table_name, schema, **kw + ) + pk_constraints = [c for c in constraints if c.contype == "p"] if not pk_constraints: - return {'constrained_columns': [], 'name': ''} + return {"constrained_columns": [], "name": ""} pk_constraint = pk_constraints[0] m = PRIMARY_KEY_RE.match(pk_constraint.condef) - colstring = m.group('columns') + colstring = m.group("columns") constrained_columns = SQL_IDENTIFIER_RE.findall(colstring) return { - 'constrained_columns': constrained_columns, - 'name': pk_constraint.conname, + "constrained_columns": constrained_columns, + "name": pk_constraint.conname, } @reflection.cache @@ -905,28 +1080,29 @@ def get_foreign_keys(self, connection, table_name, schema=None, **kw): Overrides interface :meth:`~sqlalchemy.engine.interfaces.Dialect.get_pk_constraint`. """ - constraints = self._get_redshift_constraints(connection, table_name, - schema, **kw) - fk_constraints = [c for c in constraints if c.contype == 'f'] + constraints = self._get_redshift_constraints( + connection, table_name, schema, **kw + ) + fk_constraints = [c for c in constraints if c.contype == "f"] uniques = defaultdict(lambda: defaultdict(dict)) for con in fk_constraints: uniques[con.conname]["key"] = con.conkey uniques[con.conname]["condef"] = con.condef fkeys = [] for conname, attrs in uniques.items(): - m = FOREIGN_KEY_RE.match(attrs['condef']) - colstring = m.group('referred_columns') + m = FOREIGN_KEY_RE.match(attrs["condef"]) + colstring = m.group("referred_columns") referred_columns = SQL_IDENTIFIER_RE.findall(colstring) - referred_table = m.group('referred_table') - referred_schema = m.group('referred_schema') - colstring = m.group('columns') + referred_table = m.group("referred_table") + referred_schema = m.group("referred_schema") + colstring = m.group("columns") constrained_columns = SQL_IDENTIFIER_RE.findall(colstring) fkey_d = { - 'name': conname, - 'constrained_columns': constrained_columns, - 'referred_schema': referred_schema, - 'referred_table': referred_table, - 'referred_columns': referred_columns, + "name": conname, + "constrained_columns": constrained_columns, + "referred_schema": referred_schema, + "referred_table": referred_table, + "referred_columns": referred_columns, } fkeys.append(fkey_d) return fkeys @@ -939,7 +1115,7 @@ def get_table_names(self, connection, schema=None, **kw): Overrides interface :meth:`~sqlalchemy.engine.interfaces.Dialect.get_table_names`. """ - return self._get_table_or_view_names('r', connection, schema, **kw) + return self._get_table_or_view_names("r", connection, schema, **kw) @reflection.cache def get_view_names(self, connection, schema=None, **kw): @@ -949,7 +1125,7 @@ def get_view_names(self, connection, schema=None, **kw): Overrides interface :meth:`~sqlalchemy.engine.interfaces.Dialect.get_view_names`. """ - return self._get_table_or_view_names('v', connection, schema, **kw) + return self._get_table_or_view_names("v", connection, schema, **kw) @reflection.cache def get_view_definition(self, connection, view_name, schema=None, **kw): @@ -976,25 +1152,24 @@ def get_indexes(self, connection, table_name, schema, **kw): return [] @reflection.cache - def get_unique_constraints(self, connection, table_name, - schema=None, **kw): + def get_unique_constraints(self, connection, table_name, schema=None, **kw): """ Return information about unique constraints in `table_name`. Overrides interface :meth:`~sqlalchemy.engine.interfaces.Dialect.get_unique_constraints`. """ - constraints = self._get_redshift_constraints(connection, - table_name, schema, **kw) - constraints = [c for c in constraints if c.contype == 'u'] + constraints = self._get_redshift_constraints( + connection, table_name, schema, **kw + ) + constraints = [c for c in constraints if c.contype == "u"] uniques = defaultdict(lambda: defaultdict(dict)) for con in constraints: uniques[con.conname]["key"] = con.conkey uniques[con.conname]["cols"][con.attnum] = con.attname return [ - {'name': name, - 'column_names': [uc["cols"][i] for i in uc["key"]]} + {"name": name, "column_names": [uc["cols"][i] for i in uc["key"]]} for name, uc in uniques.items() ] @@ -1007,17 +1182,16 @@ def get_table_options(self, connection, table_name, schema, **kw): Overrides interface :meth:`~sqlalchemy.engine.Inspector.get_table_options`. """ + def keyfunc(column): num = int(column.sortkey) # If sortkey is interleaved, column numbers alternate # negative values, so take abs. return abs(num) - table = self._get_redshift_relation(connection, table_name, - schema, **kw) - columns = self._get_redshift_columns(connection, table_name, - schema, **kw) - sortkey_cols = sorted([col for col in columns if col.sortkey], - key=keyfunc) + + table = self._get_redshift_relation(connection, table_name, schema, **kw) + columns = self._get_redshift_columns(connection, table_name, schema, **kw) + sortkey_cols = sorted([col for col in columns if col.sortkey], key=keyfunc) interleaved = any([int(col.sortkey) < 0 for col in sortkey_cols]) sortkey = tuple(col.name for col in sortkey_cols) interleaved_sortkey = None @@ -1027,20 +1201,20 @@ def keyfunc(column): distkeys = [col.name for col in columns if col.distkey] distkey = distkeys[0] if distkeys else None return { - 'redshift_diststyle': table.diststyle, - 'redshift_distkey': distkey, - 'redshift_sortkey': sortkey, - 'redshift_interleaved_sortkey': interleaved_sortkey, + "redshift_diststyle": table.diststyle, + "redshift_distkey": distkey, + "redshift_sortkey": sortkey, + "redshift_interleaved_sortkey": interleaved_sortkey, } def _get_table_or_view_names(self, relkind, connection, schema=None, **kw): default_schema = inspect(connection).default_schema_name if not schema: schema = default_schema - info_cache = kw.get('info_cache') - all_relations = self._get_all_relation_info(connection, - schema=schema, - info_cache=info_cache) + info_cache = kw.get("info_cache") + all_relations = self._get_all_relation_info( + connection, schema=schema, info_cache=info_cache + ) relation_names = [] for key, relation in all_relations.items(): if key.schema == schema and relation.relkind == relkind: @@ -1049,37 +1223,32 @@ def _get_table_or_view_names(self, relkind, connection, schema=None, **kw): def _get_column_info(self, *args, **kwargs): kw = kwargs.copy() - encode = kw.pop('encode', None) - if sa_version >= Version('1.3.16'): + encode = kw.pop("encode", None) + if sa_version >= Version("1.3.16"): # SQLAlchemy 1.3.16 introduced generated columns, # not supported in redshift - kw['generated'] = '' - - if sa_version < Version('1.4.0') and 'identity' in kw: - del kw['identity'] - elif sa_version >= Version('1.4.0') and 'identity' not in kw: - kw['identity'] = None - - column_info = super(RedshiftDialectMixin, self)._get_column_info( - *args, - **kw - ) - if isinstance(column_info['type'], VARCHAR): - if column_info['type'].length is None: - column_info['type'] = NullType() - if 'info' not in column_info: - column_info['info'] = {} - if encode and encode != 'none': - column_info['info']['encode'] = encode + kw["generated"] = "" + + if sa_version < Version("1.4.0") and "identity" in kw: + del kw["identity"] + elif sa_version >= Version("1.4.0") and "identity" not in kw: + kw["identity"] = None + + column_info = super(RedshiftDialectMixin, self)._get_column_info(*args, **kw) + if isinstance(column_info["type"], VARCHAR): + if column_info["type"].length is None: + column_info["type"] = NullType() + if "info" not in column_info: + column_info["info"] = {} + if encode and encode != "none": + column_info["info"]["encode"] = encode return column_info - def _get_redshift_relation(self, connection, table_name, - schema=None, **kw): - info_cache = kw.get('info_cache') - all_relations = self._get_all_relation_info(connection, - schema=schema, - table_name=table_name, - info_cache=info_cache) + def _get_redshift_relation(self, connection, table_name, schema=None, **kw): + info_cache = kw.get("info_cache") + all_relations = self._get_all_relation_info( + connection, schema=schema, table_name=table_name, info_cache=info_cache + ) key = RelationKey(table_name, schema, connection) if key not in all_relations.keys(): key = key.unquoted() @@ -1089,25 +1258,20 @@ def _get_redshift_relation(self, connection, table_name, raise sa.exc.NoSuchTableError(key) def _get_redshift_columns(self, connection, table_name, schema=None, **kw): - info_cache = kw.get('info_cache') + info_cache = kw.get("info_cache") all_schema_columns = self._get_schema_column_info( - connection, - schema=schema, - table_name=table_name, - info_cache=info_cache + connection, schema=schema, table_name=table_name, info_cache=info_cache ) key = RelationKey(table_name, schema, connection) if key not in all_schema_columns.keys(): key = key.unquoted() return all_schema_columns[key] - def _get_redshift_constraints(self, connection, table_name, - schema=None, **kw): - info_cache = kw.get('info_cache') - all_constraints = self._get_all_constraint_info(connection, - schema=schema, - table_name=table_name, - info_cache=info_cache) + def _get_redshift_constraints(self, connection, table_name, schema=None, **kw): + info_cache = kw.get("info_cache") + all_constraints = self._get_all_constraint_info( + connection, schema=schema, table_name=table_name, info_cache=info_cache + ) key = RelationKey(table_name, schema, connection) if key not in all_constraints.keys(): key = key.unquoted() @@ -1115,19 +1279,19 @@ def _get_redshift_constraints(self, connection, table_name, @reflection.cache def _get_all_relation_info(self, connection, **kw): - schema = kw.get('schema', None) + schema = kw.get("schema", None) schema_clause = ( "AND schema = '{schema}'".format(schema=schema) if schema else "" ) - table_name = kw.get('table_name', None) + table_name = kw.get("table_name", None) table_clause = ( - "AND relname = '{table}'".format( - table=table_name - ) if table_name else "" + "AND relname = '{table}'".format(table=table_name) if table_name else "" ) - result = connection.execute(sa.text(""" + result = connection.execute( + sa.text( + """ SELECT c.relkind, n.oid as "schema_oid", @@ -1165,7 +1329,9 @@ def _get_all_relation_info(self, connection, **kw): JOIN pg_catalog.pg_user u ON u.usesysid = s.esowner where 1 {schema_clause} {table_clause} ORDER BY "relkind", "schema_oid", "schema"; - """.format(schema_clause=schema_clause, table_clause=table_clause))) + """.format(schema_clause=schema_clause, table_clause=table_clause) + ) + ) relations = {} for rel in result: key = RelationKey(rel.relname, rel.schema, connection) @@ -1176,23 +1342,24 @@ def _get_all_relation_info(self, connection, **kw): # when reflecting schema for multiple tables at once. @reflection.cache def _get_schema_column_info(self, connection, **kw): - schema = kw.get('schema', None) + schema = kw.get("schema", None) schema_clause = ( "AND schema = '{schema}'".format(schema=schema) if schema else "" ) - table_name = kw.get('table_name', None) + table_name = kw.get("table_name", None) table_clause = ( - "AND table_name = '{table}'".format( - table=table_name - ) if table_name else "" + "AND table_name = '{table}'".format(table=table_name) if table_name else "" ) all_columns = defaultdict(list) - result = connection.execute(sa.text(REFLECTION_SQL.format( - schema_clause=schema_clause, - table_clause=table_clause - ))) + result = connection.execute( + sa.text( + REFLECTION_SQL.format( + schema_clause=schema_clause, table_clause=table_clause + ) + ) + ) for col in result: key = RelationKey(col.table_name, col.schema, connection) @@ -1202,19 +1369,19 @@ def _get_schema_column_info(self, connection, **kw): @reflection.cache def _get_all_constraint_info(self, connection, **kw): - schema = kw.get('schema', None) + schema = kw.get("schema", None) schema_clause = ( "AND schema = '{schema}'".format(schema=schema) if schema else "" ) - table_name = kw.get('table_name', None) + table_name = kw.get("table_name", None) table_clause = ( - "AND table_name = '{table}'".format( - table=table_name - ) if table_name else "" + "AND table_name = '{table}'".format(table=table_name) if table_name else "" ) - result = connection.execute(sa.text(""" + result = connection.execute( + sa.text( + """ SELECT n.nspname as "schema", c.relname as "table_name", @@ -1251,7 +1418,9 @@ def _get_all_constraint_info(self, connection, **kw): JOIN svv_external_schemas s ON s.schemaname = c.schemaname where 1 {schema_clause} {table_clause} ORDER BY "schema", "table_name" - """.format(schema_clause=schema_clause, table_clause=table_clause))) + """.format(schema_clause=schema_clause, table_clause=table_clause) + ) + ) all_constraints = defaultdict(list) for con in result: key = RelationKey(con.table_name, con.schema, connection) @@ -1270,6 +1439,7 @@ class Psycopg2RedshiftDialectMixin(RedshiftDialectMixin): :class:`~sqlalchemy.engine.interfaces.Dialect` and :class:`~sqlalchemy.engine.Inspector`. """ + def create_connect_args(self, *args, **kwargs): """ Build DB-API compatible connection arguments. @@ -1278,33 +1448,28 @@ def create_connect_args(self, *args, **kwargs): :meth:`~sqlalchemy.engine.interfaces.Dialect.create_connect_args`. """ default_args = { - 'sslmode': 'verify-full', - 'sslrootcert': pkg_resources.resource_filename( - __name__, - 'redshift-ca-bundle.crt' + "sslmode": "verify-full", + "sslrootcert": str( + files("sqlalchemy_redshift").joinpath("redshift-ca-bundle.crt") ), } - cargs, cparams = ( - super(Psycopg2RedshiftDialectMixin, self).create_connect_args( - *args, **kwargs - ) + cargs, cparams = super(Psycopg2RedshiftDialectMixin, self).create_connect_args( + *args, **kwargs ) default_args.update(cparams) return cargs, default_args @classmethod def dbapi(cls): - try: - return importlib.import_module(cls.driver) - except ImportError: - raise ImportError( - 'No module named {}'.format(cls.driver) - ) + # driver_module = importlib.import_module(cls.driver) + # if Version(driver_module.__version__) < Version("2.0.908"): + # cls.description_encoding = "use_encoding" + # else: + # cls.description_encoding = None + cls.description_encoding = "use_encoding" -class RedshiftDialect_psycopg2( - Psycopg2RedshiftDialectMixin, PGDialect_psycopg2 -): +class RedshiftDialect_psycopg2(Psycopg2RedshiftDialectMixin, PGDialect_psycopg2): supports_statement_cache = False @@ -1319,7 +1484,6 @@ class RedshiftDialect_psycopg2cffi( class RedshiftDialect_redshift_connector(RedshiftDialectMixin, PGDialect): - class RedshiftCompiler_redshift_connector(RedshiftCompiler, PGCompiler): def limit_clause(self, select, **kw): text = "" @@ -1342,6 +1506,7 @@ def visit_mod_binary(self, binary, operator, **kw): def post_process_text(self, text): from sqlalchemy import util + if "%%" in text: util.warn( "The SQLAlchemy postgresql dialect " @@ -1355,7 +1520,7 @@ def pre_exec(self): if not self.compiled: return - driver = 'redshift_connector' + driver = "redshift_connector" supports_unicode_statements = True @@ -1370,9 +1535,9 @@ def pre_exec(self): use_setinputsizes = False # not implemented in redshift_connector def __init__(self, client_encoding=None, **kwargs): - super( - RedshiftDialect_redshift_connector, self - ).__init__(client_encoding=client_encoding, **kwargs) + super(RedshiftDialect_redshift_connector, self).__init__( + client_encoding=client_encoding, **kwargs + ) self.client_encoding = client_encoding @classmethod @@ -1381,7 +1546,7 @@ def dbapi(cls): driver_module = importlib.import_module(cls.driver) # Starting v2.0.908 driver converts description column names to str - if Version(driver_module.__version__) < Version('2.0.908'): + if Version(driver_module.__version__) < Version("2.0.908"): cls.description_encoding = "use_encoding" else: cls.description_encoding = None @@ -1389,8 +1554,8 @@ def dbapi(cls): return driver_module except ImportError: raise ImportError( - 'No module named redshift_connector. Please install ' - 'redshift_connector to use this sqlalchemy dialect.' + "No module named redshift_connector. Please install " + "redshift_connector to use this sqlalchemy dialect." ) def set_client_encoding(self, connection, client_encoding): @@ -1427,9 +1592,9 @@ def set_isolation_level(self, connection, level): connection.autocommit = True else: connection.autocommit = False - super( - RedshiftDialect_redshift_connector, self - ).set_isolation_level(connection, level) + super(RedshiftDialect_redshift_connector, self).set_isolation_level( + connection, level + ) def on_connect(self): fns = [] @@ -1437,6 +1602,7 @@ def on_connect(self): def on_connect(conn): from sqlalchemy import util from sqlalchemy.sql.elements import quoted_name + conn.py_types[quoted_name] = conn.py_types[util.text_type] fns.append(on_connect) @@ -1473,25 +1639,23 @@ def create_connect_args(self, *args, **kwargs): :meth:`~sqlalchemy.engine.interfaces.Dialect.create_connect_args`. """ default_args = { - 'sslmode': 'verify-full', - 'ssl': True, - 'application_name': 'sqlalchemy-redshift' + "sslmode": "verify-full", + "ssl": True, + "application_name": "sqlalchemy-redshift", } cargs, cparams = super(RedshiftDialectMixin, self).create_connect_args( *args, **kwargs ) # set client_encoding so it is picked up by on_connect(), as # redshift_connector does not have client_encoding connection parameter - self.client_encoding = cparams.pop( - 'client_encoding', self.client_encoding - ) + self.client_encoding = cparams.pop("client_encoding", self.client_encoding) - if 'port' in cparams: - cparams['port'] = int(cparams['port']) + if "port" in cparams: + cparams["port"] = int(cparams["port"]) - if 'username' in cparams: - cparams['user'] = cparams['username'] - del cparams['username'] + if "username" in cparams: + cparams["user"] = cparams["username"] + del cparams["username"] default_args.update(cparams) return cargs, default_args @@ -1513,7 +1677,7 @@ def gen_columns_from_children(root): yield root -@compiles(Delete, 'redshift') +@compiles(Delete, "redshift") def visit_delete_stmt(element, compiler, **kwargs): """ Adds redshift-dialect specific compilation rule for the @@ -1569,8 +1733,8 @@ def visit_delete_stmt(element, compiler, **kwargs): """ # Set empty strings for the default where clause and using clause - whereclause = '' - usingclause = '' + whereclause = "" + usingclause = "" # determine if the delete query needs a ``USING`` injected # by inspecting the whereclause's children & their children... @@ -1581,15 +1745,15 @@ def visit_delete_stmt(element, compiler, **kwargs): # which they first appear in the where clause. delete_stmt_table = compiler.process(element.table, asfrom=True, **kwargs) - if sa_version >= Version('1.4.0'): + if sa_version >= Version("1.4.0"): if element.whereclause is not None: clause = compiler.process(element.whereclause, **kwargs) if clause: - whereclause = ' WHERE {clause}'.format(clause=clause) + whereclause = " WHERE {clause}".format(clause=clause) else: whereclause_tuple = element.get_children() if whereclause_tuple: - whereclause = ' WHERE {clause}'.format( + whereclause = " WHERE {clause}".format( clause=compiler.process(*whereclause_tuple, **kwargs) ) @@ -1598,15 +1762,11 @@ def visit_delete_stmt(element, compiler, **kwargs): whereclause_columns = gen_columns_from_children(element) for col in whereclause_columns: table = compiler.process(col.table, asfrom=True, **kwargs) - if table != delete_stmt_table and \ - table not in usingclause_tables: + if table != delete_stmt_table and table not in usingclause_tables: usingclause_tables.append(table) if usingclause_tables: - usingclause = ' USING {clause}'.format( - clause=', '.join(usingclause_tables) - ) + usingclause = " USING {clause}".format(clause=", ".join(usingclause_tables)) - return 'DELETE FROM {table}{using}{where}'.format( - table=delete_stmt_table, - using=usingclause, - where=whereclause) + return "DELETE FROM {table}{using}{where}".format( + table=delete_stmt_table, using=usingclause, where=whereclause + ) diff --git a/tests/test_reflection_views.py b/tests/test_reflection_views.py index d15470f6..5921b66e 100644 --- a/tests/test_reflection_views.py +++ b/tests/test_reflection_views.py @@ -1,12 +1,11 @@ +from rs_sqla_test_utils.utils import clean, compile_query +import sqlalchemy as sa from sqlalchemy import MetaData, Table, inspect from sqlalchemy.schema import CreateTable -import sqlalchemy as sa -from rs_sqla_test_utils.utils import clean, compile_query def table_to_ddl(engine, table): - return str(CreateTable(table) - .compile(engine)) + return str(CreateTable(table).compile(engine)) def test_view_reflection(redshift_engine): @@ -20,24 +19,24 @@ def test_view_reflection(redshift_engine): conn.execute(sa.text(view_ddl)) conn.execute(sa.text("COMMIT")) insp = inspect(redshift_engine) - view_definition = insp.get_view_definition('my_view') + view_definition = insp.get_view_definition("my_view") assert clean( compile_query(view_definition, redshift_engine.dialect) ) == clean(view_query) - view = Table('my_view', MetaData(), - autoload=True, autoload_with=redshift_engine) - assert(len(view.columns) == 2) + view = Table( + "my_view", MetaData(), autoload=True, autoload_with=redshift_engine + ) + assert len(view.columns) == 2 finally: - conn.execute(sa.text('DROP TABLE IF EXISTS my_table CASCADE')) - conn.execute(sa.text('DROP VIEW IF EXISTS my_view CASCADE')) + conn.execute(sa.text("DROP TABLE IF EXISTS my_table CASCADE")) + conn.execute(sa.text("DROP VIEW IF EXISTS my_view CASCADE")) conn.execute(sa.text("COMMIT")) def test_late_binding_view_reflection(redshift_engine): table_ddl = "CREATE TABLE my_table (col1 INTEGER, col2 INTEGER)" view_query = "SELECT my_table.col1, my_table.col2 FROM public.my_table" - view_ddl = ("CREATE VIEW my_late_view AS " - "%s WITH NO SCHEMA BINDING" % view_query) + view_ddl = "CREATE VIEW my_late_view AS %s WITH NO SCHEMA BINDING" % view_query with redshift_engine.connect() as conn: try: @@ -45,16 +44,17 @@ def test_late_binding_view_reflection(redshift_engine): conn.execute(sa.text(view_ddl)) conn.execute(sa.text("COMMIT")) insp = inspect(redshift_engine) - view_definition = insp.get_view_definition('my_late_view') + view_definition = insp.get_view_definition("my_late_view") # Redshift returns the entire DDL for late binding views. assert clean( compile_query(view_definition, redshift_engine.dialect) ) == clean(view_ddl) - view = Table('my_late_view', MetaData(), - autoload=True, autoload_with=redshift_engine) - assert(len(view.columns) == 2) + view = Table( + "my_late_view", MetaData(), autoload=True, autoload_with=redshift_engine + ) + assert len(view.columns) == 2 finally: - conn.execute(sa.text('DROP TABLE IF EXISTS my_table CASCADE')) - conn.execute(sa.text('DROP VIEW IF EXISTS my_late_view CASCADE')) - conn.execute(sa.text('COMMIT')) + conn.execute(sa.text("DROP TABLE IF EXISTS my_table CASCADE")) + conn.execute(sa.text("DROP VIEW IF EXISTS my_late_view CASCADE")) + conn.execute(sa.text("COMMIT")) diff --git a/tox.ini b/tox.ini index b24a463b..eeaaad9c 100644 --- a/tox.ini +++ b/tox.ini @@ -1,9 +1,6 @@ [tox] envlist = - py39-pg28-sa13 - py39-pg28-sa14 - py310-pg28-sa13 - py310-pg28-sa14 + py310-pg29-sa14 lint docs @@ -11,22 +8,19 @@ envlist = commands = pytest {posargs} --dbdriver psycopg2 --dbdriver psycopg2cffi --dbdriver redshift_connector passenv = PGPASSWORD,REDSHIFT_USERNAME,REDSHIFT_HOST,REDSHIFT_PORT,REDSHIFT_DATABASE,REDSHIFT_IAM_ROLE_ARN,AWS_ACCOUNT_ID,REDSHIFT_IAM_ROLE_NAME deps = - sa13: sqlalchemy==1.3.24 sa14: sqlalchemy==1.4.15 - pg28: psycopg2==2.8.6 pg29: psycopg2==2.9.5 - alembic==1.9.2 - packaging==20.4 - psycopg2cffi==2.8.1 - pytest==7.2.1 - requests==2.25.0 - redshift_connector==2.0.907 - requests==2.25.0 + alembic==1.13.1 + packaging==24.2 + psycopg2cffi==2.9.0 + pytest==8.3.4 + requests==2.32.3 + redshift_connector==2.1.4 [testenv:lint] deps = - flake8==4.0.1 - psycopg2 + flake8==7.1.1 + psycopg2-binary redshift_connector commands=flake8 sqlalchemy_redshift tests From 8658e0e32161911f896f45dab183bb48b15a485c Mon Sep 17 00:00:00 2001 From: "A. Fox" Date: Tue, 4 Aug 2026 11:11:38 -0700 Subject: [PATCH 10/12] fix: don't use description_encoding --- sqlalchemy_redshift/dialect.py | 27 ++++++--------------------- 1 file changed, 6 insertions(+), 21 deletions(-) diff --git a/sqlalchemy_redshift/dialect.py b/sqlalchemy_redshift/dialect.py index bba4f227..72fa8682 100644 --- a/sqlalchemy_redshift/dialect.py +++ b/sqlalchemy_redshift/dialect.py @@ -1461,12 +1461,9 @@ def create_connect_args(self, *args, **kwargs): @classmethod def dbapi(cls): - # driver_module = importlib.import_module(cls.driver) - # if Version(driver_module.__version__) < Version("2.0.908"): - # cls.description_encoding = "use_encoding" - # else: - # cls.description_encoding = None - cls.description_encoding = "use_encoding" + driver_module = importlib.import_module(cls.driver) + cls.description_encoding = None + return driver_module class RedshiftDialect_psycopg2(Psycopg2RedshiftDialectMixin, PGDialect_psycopg2): @@ -1542,21 +1539,9 @@ def __init__(self, client_encoding=None, **kwargs): @classmethod def dbapi(cls): - try: - driver_module = importlib.import_module(cls.driver) - - # Starting v2.0.908 driver converts description column names to str - if Version(driver_module.__version__) < Version("2.0.908"): - cls.description_encoding = "use_encoding" - else: - cls.description_encoding = None - - return driver_module - except ImportError: - raise ImportError( - "No module named redshift_connector. Please install " - "redshift_connector to use this sqlalchemy dialect." - ) + driver_module = importlib.import_module(cls.driver) + cls.description_encoding = None + return driver_module def set_client_encoding(self, connection, client_encoding): """ From 0161beb2055edad646fa60d7527ab03c5f522b1c Mon Sep 17 00:00:00 2001 From: "A. Fox" Date: Tue, 4 Aug 2026 11:43:23 -0700 Subject: [PATCH 11/12] update sqlalchemy supported version --- tests/test_default_ssl.py | 23 +++++++++++++---------- tox.ini | 5 +++-- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/tests/test_default_ssl.py b/tests/test_default_ssl.py index f5586d30..9f486b8c 100644 --- a/tests/test_default_ssl.py +++ b/tests/test_default_ssl.py @@ -1,26 +1,29 @@ +# from pkg_resources import resource_filename +from importlib.resources import files + import sqlalchemy as sa -from pkg_resources import resource_filename + from sqlalchemy_redshift.dialect import ( - Psycopg2RedshiftDialectMixin, RedshiftDialect_redshift_connector + Psycopg2RedshiftDialectMixin, + RedshiftDialect_redshift_connector, ) - -CERT_PATH = resource_filename("sqlalchemy_redshift", "redshift-ca-bundle.crt") +CERT_PATH = str(files("sqlalchemy_redshift").joinpath("redshift-ca-bundle.crt")) def test_ssl_args(redshift_dialect_flavor): - engine = sa.create_engine('{}://test'.format(redshift_dialect_flavor)) + engine = sa.create_engine("{}://test".format(redshift_dialect_flavor)) dialect = engine.dialect url = engine.url cargs, cparams = dialect.create_connect_args(url) assert cargs == [] - assert cparams.pop('host') == 'test' - assert cparams.pop('sslmode') == 'verify-full' + assert cparams.pop("host") == "test" + assert cparams.pop("sslmode") == "verify-full" if isinstance(dialect, Psycopg2RedshiftDialectMixin): - assert cparams.pop('sslrootcert') == CERT_PATH + assert cparams.pop("sslrootcert") == CERT_PATH elif isinstance(dialect, RedshiftDialect_redshift_connector): - assert cparams.pop('ssl') is True - assert cparams.pop('application_name') == 'sqlalchemy-redshift' + assert cparams.pop("ssl") is True + assert cparams.pop("application_name") == "sqlalchemy-redshift" assert cparams == {} diff --git a/tox.ini b/tox.ini index eeaaad9c..3966770d 100644 --- a/tox.ini +++ b/tox.ini @@ -1,6 +1,7 @@ [tox] envlist = py310-pg29-sa14 + py313-pg29-sa14 lint docs @@ -8,8 +9,8 @@ envlist = commands = pytest {posargs} --dbdriver psycopg2 --dbdriver psycopg2cffi --dbdriver redshift_connector passenv = PGPASSWORD,REDSHIFT_USERNAME,REDSHIFT_HOST,REDSHIFT_PORT,REDSHIFT_DATABASE,REDSHIFT_IAM_ROLE_ARN,AWS_ACCOUNT_ID,REDSHIFT_IAM_ROLE_NAME deps = - sa14: sqlalchemy==1.4.15 - pg29: psycopg2==2.9.5 + sa14: sqlalchemy==1.4.54 + pg29: psycopg2==2.9.12 alembic==1.13.1 packaging==24.2 psycopg2cffi==2.9.0 From 5208feb52cc5f56cf70a132e2667ace410e00831 Mon Sep 17 00:00:00 2001 From: "A. Fox" Date: Tue, 4 Aug 2026 12:28:13 -0700 Subject: [PATCH 12/12] fix: remove commented import --- tests/test_default_ssl.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_default_ssl.py b/tests/test_default_ssl.py index 9f486b8c..291be37d 100644 --- a/tests/test_default_ssl.py +++ b/tests/test_default_ssl.py @@ -1,4 +1,3 @@ -# from pkg_resources import resource_filename from importlib.resources import files import sqlalchemy as sa