From b7abd41478f30a8875004ce19aceffe015d0f797 Mon Sep 17 00:00:00 2001 From: Gord Thompson Date: Fri, 10 Apr 2026 15:49:51 -0600 Subject: [PATCH 1/4] re-work get_multi_columns Fixes: #297 --- .github/workflows/ci.yml | 8 +- CHANGES.md | 1 + dev-requirements.txt | 43 ++++-- sqlalchemy_cockroachdb/base.py | 231 ++++++++++++---------------- test-requirements.txt | 18 ++- test/test_column_reflect.py | 53 ++++++- test/test_introspection.py | 92 ++++-------- test/test_suite_sqlalchemy.py | 265 +++++++++++++++++---------------- 8 files changed, 363 insertions(+), 348 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1e55cc8..bb6e875 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,8 +54,8 @@ jobs: TOXENV: py310 TOX_VERSION: 4.34.1 steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 with: python-version: '3.10' - name: Start CockroachDB @@ -77,8 +77,8 @@ jobs: TOXENV: py310 TOX_VERSION: 4.34.1 steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 with: python-version: '3.10' - name: Install testrunner diff --git a/CHANGES.md b/CHANGES.md index 7970154..abede13 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -5,6 +5,7 @@ Unreleased - Fix reflection of TIMESTAMPTZ columns (#276), thanks to @nvachhar - Fix reflection of JSONB columns (#277) - Fix compatibility issues with Alembic 1.18 (via SQLA 2.0.47) +- Fix get_columns() to include identity column info (#297) - Update minimum Python version to 3.10 diff --git a/dev-requirements.txt b/dev-requirements.txt index a8d9c9d..aa31d0f 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -1,12 +1,14 @@ -cachetools==7.0.1 +backports-tarfile==1.2.0 + # via jaraco-context +cachetools==7.0.5 # via tox -certifi==2026.1.4 +certifi==2026.2.25 # via requests cffi==2.0.0 # via cryptography -chardet==6.0.0.post1 +chardet==7.4.1 # via tox -charset-normalizer==3.4.4 +charset-normalizer==3.4.7 # via requests colorama==0.4.6 # via tox @@ -16,17 +18,20 @@ distlib==0.4.0 # via virtualenv docutils==0.22.4 # via readme-renderer -filelock==3.24.3 +filelock==3.25.2 # via + # python-discovery # tox # virtualenv id==1.6.1 # via twine idna==3.11 # via requests +importlib-metadata==9.0.0 + # via keyring jaraco-classes==3.4.0 # via keyring -jaraco-context==6.1.0 +jaraco-context==6.1.2 # via keyring jaraco-functools==4.4.0 # via keyring @@ -40,19 +45,20 @@ markdown-it-py==4.0.0 # via rich mdurl==0.1.2 # via markdown-it-py -more-itertools==10.8.0 +more-itertools==11.0.2 # via # jaraco-classes # jaraco-functools -nh3==0.3.3 +nh3==0.3.4 # via readme-renderer packaging==26.0 # via # pyproject-api # tox # twine -platformdirs==4.9.2 +platformdirs==4.9.6 # via + # python-discovery # tox # virtualenv pluggy==1.6.0 @@ -65,9 +71,11 @@ pygments==2.20.0 # rich pyproject-api==1.10.0 # via tox +python-discovery==1.2.2 + # via virtualenv readme-renderer==44.0 # via twine -requests==2.33.0 +requests==2.33.1 # via # requests-toolbelt # twine @@ -75,18 +83,29 @@ requests-toolbelt==1.0.0 # via twine rfc3986==2.0.0 # via twine -rich==14.3.3 +rich==15.0.0 # via twine secretstorage==3.5.0 # via keyring +tomli==2.4.1 + # via + # pyproject-api + # tox tox==4.34.1 # via -r dev-requirements.in twine==6.2.0 # via -r dev-requirements.in +typing-extensions==4.15.0 + # via + # cryptography + # tox + # virtualenv urllib3==2.6.3 # via # id # requests # twine -virtualenv==20.39.0 +virtualenv==21.2.1 # via tox +zipp==3.23.0 + # via importlib-metadata diff --git a/sqlalchemy_cockroachdb/base.py b/sqlalchemy_cockroachdb/base.py index 8ca0bec..ac019a8 100644 --- a/sqlalchemy_cockroachdb/base.py +++ b/sqlalchemy_cockroachdb/base.py @@ -1,21 +1,30 @@ import collections -import re import threading from sqlalchemy import text +from sqlalchemy import ARRAY +from sqlalchemy import BIGINT +from sqlalchemy import BLOB +from sqlalchemy import DECIMAL +from sqlalchemy import DOUBLE_PRECISION +from sqlalchemy import FLOAT +from sqlalchemy import INTEGER +from sqlalchemy import NUMERIC +from sqlalchemy import REAL +from sqlalchemy import SMALLINT +from sqlalchemy import TEXT +from sqlalchemy import VARCHAR from sqlalchemy.dialects.postgresql.base import PGDialect -from sqlalchemy.dialects.postgresql import ARRAY +from sqlalchemy.dialects.postgresql import BYTEA from sqlalchemy.dialects.postgresql import INET -from sqlalchemy.dialects.postgresql import UUID from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.dialects.postgresql import UUID from sqlalchemy.ext.compiler import compiles -from sqlalchemy.util import warn import sqlalchemy.types as sqltypes from .stmt_compiler import CockroachCompiler, CockroachIdentifierPreparer from .ddl_compiler import CockroachDDLCompiler - # Map type names (as returned by information_schema) to sqlalchemy type # objects. # @@ -151,7 +160,7 @@ def _get_server_version_info(self, conn): # PGDialect expects a postgres server version number here, # although we've overridden most of the places where it's # used. - return (9, 5, 0) + return (12, 0, 0) def get_table_names(self, conn, schema=None, **kw): # Upstream implementation needs correlated subqueries. @@ -174,126 +183,88 @@ def has_table(self, conn, table, schema=None, info_cache=None): return any(t == table for t in self.get_table_names(conn, schema=schema)) def get_multi_columns(self, connection, schema, filter_names, scope, kind, **kw): - if not filter_names: - filter_names = self.get_table_names(connection, schema) - return { - (schema, table_name): self.get_columns(connection, table_name, schema, **kw) - for table_name in filter_names - } - - # The upstream implementations of the reflection functions below depend on - # correlated subqueries which are not yet supported. - def get_columns(self, conn, table_name, schema=None, **kw): _include_hidden = kw.get("include_hidden", False) - if not self._is_v191plus: - # v2.x does not have is_generated or generation_expression - sql = ( - "SELECT column_name, data_type, is_nullable::bool, column_default," - "numeric_precision, numeric_scale, character_maximum_length, " - "NULL AS is_generated, NULL AS generation_expression, is_hidden::bool," - "column_comment AS comment " - "FROM information_schema.columns " - "WHERE table_schema = :table_schema AND table_name = :table_name " - ) - sql += "" if _include_hidden else "AND NOT is_hidden::bool" - rows = conn.execute( - text(sql), - {"table_schema": schema or self.default_schema_name, "table_name": table_name}, - ) - else: - # v19.1 or later. Information schema columns are all usable. - sql = ( - "SELECT column_name, data_type, is_nullable::bool, column_default, " - "numeric_precision, numeric_scale, character_maximum_length, " - "CASE is_generated WHEN 'ALWAYS' THEN true WHEN 'NEVER' THEN false " - "ELSE is_generated::bool END AS is_generated, " - "generation_expression, is_hidden::bool, crdb_sql_type, column_comment AS comment " - "FROM information_schema.columns " - "WHERE table_schema = :table_schema AND table_name = :table_name " - ) - sql += "" if _include_hidden else "AND NOT is_hidden::bool" - rows = conn.execute( - text(sql), - {"table_schema": schema or self.default_schema_name, "table_name": table_name}, - ) - - res = [] - for row in rows: - name, type_str, nullable, default = row[:4] - if type_str == "ARRAY": - is_array = True - type_str, _ = row.crdb_sql_type.split("[", maxsplit=1) - else: - is_array = False - # When there are type parameters, attach them to the - # returned type object. - m = re.match(r"^(\w+(?: \w+)*)(?:\(([0-9, ]*)\))?$", type_str) - if m is None: - warn("Could not parse type name '%s'" % type_str) - typ = sqltypes.NULLTYPE - else: - type_name, type_args = m.groups() - try: - type_class = _type_map[type_name.lower()] - except KeyError: - warn(f"Did not recognize type '{type_name}' of column '{name}'") - type_class = sqltypes.NULLTYPE - if type_args: - typ = type_class(*[int(s.strip()) for s in type_args.split(",")]) - elif type_class is sqltypes.DECIMAL: - typ = type_class( - precision=row.numeric_precision, - scale=row.numeric_scale, - ) - elif type_class is sqltypes.VARCHAR or type_class is sqltypes.CHAR: - typ = type_class(length=row.character_maximum_length) - else: - typ = type_class - if row.is_generated: - # Currently, all computed columns are persisted. - computed = dict(sqltext=row.generation_expression, persisted=True) - default = None - else: - computed = None - # Check if a sequence is being used and adjust the default value. - autoincrement = False - if default is not None: - nextval_match = re.search(r"""(nextval\(')([^']+)('.*$)""", default) - unique_rowid_match = re.search(r"""unique_rowid\(""", default) - if nextval_match is not None or unique_rowid_match is not None: - if issubclass(type_class, sqltypes.Integer): - autoincrement = True - # the default is related to a Sequence - sch = schema - if ( - nextval_match is not None - and "." not in nextval_match.group(2) - and sch is not None - ): - # unconditionally quote the schema name. this could - # later be enhanced to obey quoting rules / - # "quote schema" - default = ( - nextval_match.group(1) - + ('"%s"' % sch) - + "." - + nextval_match.group(2) - + nextval_match.group(3) + multi_columns = super().get_multi_columns( + connection, schema, filter_names, scope, kind, **kw + ) + to_return = [] + if multi_columns: + current = connection.execute( + text("select current_database() as db, current_schema() as schema") + ).one() + for table, columns in multi_columns: + if table not in [ + (None, "geography_columns"), + (None, "geometry_columns"), + (None, "spatial_ref_sys"), + ]: + table_columns = ( + connection.execute( + text( + "select column_name, is_hidden::bool " + "from information_schema.columns " + "where table_catalog = :tc " + "and table_schema = :ts and table_name = :tn" + ), + dict(tc=current.db, ts=table[0] or current.schema, tn=table[1]), ) - - column_info = dict( - name=name, - type=ARRAY(typ) if is_array else typ, - nullable=nullable, - default=default, - autoincrement=autoincrement, - is_hidden=row.is_hidden, - comment=row.comment, - ) - if computed is not None: - column_info["computed"] = computed - res.append(column_info) - return res + .mappings() + .all() + ) + is_hidden = {x["column_name"]: x["is_hidden"] for x in table_columns} + for col in columns: + if is_hidden[col["name"]] and not _include_hidden: + columns.remove(col) + else: + col["is_hidden"] = is_hidden[col["name"]] + if col["default"] == "unique_rowid()": + col["autoincrement"] = True + if isinstance(col["type"], BIGINT): + col["type"] = INTEGER + elif isinstance(col["type"], ARRAY) and isinstance( + col["type"].item_type, BIGINT + ): + col["type"].item_type = INTEGER() + elif isinstance(col["type"], BYTEA): + col["type"] = BLOB + elif isinstance(col["type"], ARRAY) and isinstance( + col["type"].item_type, BYTEA + ): + col["type"].item_type = BLOB() + elif isinstance(col["type"], DOUBLE_PRECISION): + col["type"] = FLOAT + elif isinstance(col["type"], ARRAY) and isinstance( + col["type"].item_type, DOUBLE_PRECISION + ): + col["type"].item_type = FLOAT() + elif isinstance(col["type"], NUMERIC): + col["type"] = DECIMAL(col["type"].precision, col["type"].scale) + elif isinstance(col["type"], ARRAY) and isinstance( + col["type"].item_type, NUMERIC + ): + col["type"].item_type = DECIMAL( + col["type"].item_type.precision, col["type"].item_type.scale + ) + elif isinstance(col["type"], REAL): + col["type"] = FLOAT + elif isinstance(col["type"], ARRAY) and isinstance( + col["type"].item_type, REAL + ): + col["type"].item_type = FLOAT() + elif isinstance(col["type"], SMALLINT): + col["type"] = INTEGER + elif isinstance(col["type"], ARRAY) and isinstance( + col["type"].item_type, SMALLINT + ): + col["type"].item_type = INTEGER() + elif isinstance(col["type"], TEXT): + col["type"] = VARCHAR + elif isinstance(col["type"], ARRAY) and isinstance( + col["type"].item_type, TEXT + ): + col["type"].item_type = VARCHAR() + to_return.append((table, columns)) + return to_return def get_indexes(self, conn, table_name, schema=None, **kw): if self._is_v192plus: @@ -347,12 +318,8 @@ def get_indexes(self, conn, table_name, schema=None, **kw): ) return result - def get_multi_indexes( - self, connection, schema, filter_names, scope, kind, **kw - ): - result = super().get_multi_indexes( - connection, schema, filter_names, scope, kind, **kw - ) + def get_multi_indexes(self, connection, schema, filter_names, scope, kind, **kw): + result = super().get_multi_indexes(connection, schema, filter_names, scope, kind, **kw) if schema is None: result = dict(result) for k in [ @@ -417,9 +384,7 @@ def get_unique_constraints(self, conn, table_name, schema=None, **kw): res.append(index) return res - def get_multi_check_constraints( - self, connection, schema, filter_names, scope, kind, **kw - ): + def get_multi_check_constraints(self, connection, schema, filter_names, scope, kind, **kw): result = super().get_multi_check_constraints( connection, schema, filter_names, scope, kind, **kw ) diff --git a/test-requirements.txt b/test-requirements.txt index 0f59a14..be94653 100644 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -1,10 +1,14 @@ alembic==1.18.4 # via -r test-requirements.in +async-timeout==5.0.1 + # via asyncpg asyncpg==0.31.0 # via -r test-requirements.in +exceptiongroup==1.3.1 + # via pytest futures==3.0.5 # via -r test-requirements.in -greenlet==3.3.2 +greenlet==3.4.0 # via sqlalchemy iniconfig==2.3.0 # via pytest @@ -14,7 +18,7 @@ markupsafe==3.0.3 # via mako mock==5.2.0 # via -r test-requirements.in -more-itertools==10.8.0 +more-itertools==11.0.2 # via -r test-requirements.in packaging==26.0 # via pytest @@ -26,13 +30,19 @@ psycopg2==2.9.11 # via -r test-requirements.in pygments==2.20.0 # via pytest -pytest==9.0.2 +pytest==9.0.3 # via -r test-requirements.in -sqlalchemy==2.0.47 +sqlalchemy==2.0.49 # via # -r test-requirements.in # alembic +tomli==2.4.1 + # via + # alembic + # pytest typing-extensions==4.15.0 # via # alembic + # exceptiongroup + # psycopg # sqlalchemy diff --git a/test/test_column_reflect.py b/test/test_column_reflect.py index 4853cdb..33df0cf 100644 --- a/test/test_column_reflect.py +++ b/test/test_column_reflect.py @@ -1,4 +1,14 @@ -from sqlalchemy import MetaData, Table, Column, Integer, String, testing, inspect +from sqlalchemy import ( + MetaData, + Table, + Column, + Integer, + String, + testing, + inspect, + BigInteger, + Identity, +) from sqlalchemy.testing import fixtures, eq_ meta = MetaData() @@ -16,6 +26,13 @@ Column("txt", String), ) +with_identity = Table( + "with_identity", + meta, + Column("id", BigInteger, Identity(), primary_key=True), + Column("txt", String), +) + class ReflectHiddenColumnsTest(fixtures.TestBase): __requires__ = ("sync_driver",) @@ -96,3 +113,37 @@ def test_reflect_hidden_columns(self): }, ], ) + + def test_reflect_identity(self): + eq_( + self._get_col_info("with_identity"), + [ + { + "autoincrement": True, + "comment": None, + "default": "nextval('public.with_identity_id_seq'::REGCLASS)", + "identity": { + "always": False, + "cache": 1, + "cycle": False, + "increment": 1, + "maxvalue": 9223372036854775807, + "minvalue": 1, + "start": 1, + }, + "is_hidden": False, + "name": "id", + "nullable": False, + "type": "INTEGER", + }, + { + "autoincrement": False, + "comment": None, + "default": None, + "is_hidden": False, + "name": "txt", + "nullable": True, + "type": "VARCHAR", + }, + ], + ) diff --git a/test/test_introspection.py b/test/test_introspection.py index b7caba8..9e515b8 100644 --- a/test/test_introspection.py +++ b/test/test_introspection.py @@ -1,5 +1,3 @@ -import contextlib - from sqlalchemy import ( Table, Column, @@ -14,6 +12,7 @@ import sqlalchemy.types as sqltypes from sqlalchemy.testing import fixtures from sqlalchemy.dialects.postgresql import INET +from sqlalchemy.dialects.postgresql import INTERVAL from sqlalchemy.dialects.postgresql import UUID meta = MetaData() @@ -100,24 +99,42 @@ def test_array(self): self._test("timestamp[]", sqltypes.ARRAY, sqltypes.TIMESTAMP) self._test("varchar(10)[]", sqltypes.ARRAY, sqltypes.VARCHAR) + def test_blob(self): + for t in ["blob", "bytea", "bytes"]: + self._test(t, sqltypes.BLOB) + def test_boolean(self): for t in ["bool", "boolean"]: self._test(t, sqltypes.BOOLEAN) - def test_int(self): - for t in ["bigint", "int", "int2", "int4", "int64", "int8", "integer", "smallint"]: - self._test(t, sqltypes.INT) + def test_char(self): + for t in ["char", "character"]: + self._test(t, sqltypes.CHAR) - def test_float(self): - for t in ["double precision", "float", "float4", "float8", "real"]: - self._test(t, sqltypes.FLOAT) + def test_date(self): + self._test("date", sqltypes.DATE) def test_decimal(self): for t in ["dec", "decimal", "numeric"]: self._test(t, sqltypes.DECIMAL) - def test_date(self): - self._test("date", sqltypes.DATE) + def test_float(self): + for t in ["double precision", "float", "float4", "float8", "real"]: + self._test(t, sqltypes.FLOAT) + + def test_inet(self): + self._test("inet", INET) + + def test_int(self): + for t in ["bigint", "int", "int2", "int4", "int64", "int8", "integer", "smallint"]: + self._test(t, sqltypes.INT) + + def test_interval(self): + self._test("interval", INTERVAL) + + def test_json(self): + for t in ["json", "jsonb"]: + self._test(t, sqltypes.JSON) def test_time(self): for t in ["time", "time without time zone"]: @@ -133,16 +150,8 @@ def test_timestamp(self): for t in types: self._test(t, sqltypes.TIMESTAMP) - def test_interval(self): - self._test("interval", sqltypes.Interval) - - def test_char(self): - types = [ - "char", - "character", - ] - for t in types: - self._test(t, sqltypes.CHAR) + def test_uuid(self): + self._test("uuid", UUID) def test_varchar(self): types = [ @@ -154,46 +163,3 @@ def test_varchar(self): ] for t in types: self._test(t, sqltypes.VARCHAR) - - def test_blob(self): - for t in ["blob", "bytea", "bytes"]: - self._test(t, sqltypes.BLOB) - - def test_json(self): - for t in ["json", "jsonb"]: - self._test(t, sqltypes.JSON) - - def test_uuid(self): - self._test("uuid", UUID) - - def test_inet(self): - self._test("inet", INET) - - -class UnknownTypeTest(fixtures.TestBase): - __requires__ = ("sync_driver",) - - def setup_method(self): - with testing.db.begin() as conn: - conn.execute(text("CREATE TABLE t2 (c bool)")) - - def teardown_method(self): - with testing.db.begin() as conn: - conn.execute(text("DROP TABLE t2")) - - @testing.expect_warnings("Did not recognize type 'boolean'") - def test_unknown_type(self): - @contextlib.contextmanager - def make_bool_unknown(): - import sqlalchemy_cockroachdb - - t = sqlalchemy_cockroachdb.base._type_map.pop("bool") - sqlalchemy_cockroachdb.base._type_map.pop("boolean") - yield - sqlalchemy_cockroachdb.base._type_map["bool"] = t - sqlalchemy_cockroachdb.base._type_map["boolean"] = t - - with make_bool_unknown(): - meta2 = MetaData() - t = Table("t2", meta2, autoload_with=testing.db) - assert t.c["c"].type == sqltypes.NULLTYPE diff --git a/test/test_suite_sqlalchemy.py b/test/test_suite_sqlalchemy.py index b96f9e2..56edab9 100644 --- a/test/test_suite_sqlalchemy.py +++ b/test/test_suite_sqlalchemy.py @@ -1,4 +1,4 @@ -from sqlalchemy import FLOAT, INTEGER, VARCHAR +from sqlalchemy import INTEGER, VARCHAR, CHAR, DOUBLE_PRECISION from sqlalchemy.testing import skip from sqlalchemy.testing.suite import * # noqa from sqlalchemy.testing.suite import ( @@ -33,189 +33,189 @@ def test_get_multi_columns(self): insp = inspect(config.db) actual = insp.get_multi_columns() expected = { - (None, "users"): [ + (None, "comment_test"): [ { - "name": "user_id", - "type": INTEGER(), - "nullable": False, - "default": "unique_rowid()", "autoincrement": True, + "comment": "id comment", + "default": "unique_rowid()", "is_hidden": False, - "comment": None, + "name": "id", + "nullable": False, + "type": INTEGER(), }, { - "name": "test1", - "type": VARCHAR(length=5), - "nullable": False, - "default": None, "autoincrement": False, + "comment": "data % comment", + "default": None, "is_hidden": False, - "comment": None, + "name": "data", + "nullable": True, + "type": VARCHAR(length=20), }, { - "name": "test2", - "type": FLOAT(), - "nullable": False, - "default": None, "autoincrement": False, + "comment": "Comment types type speedily ' \" \ '' Fun!", # noqa + "default": None, "is_hidden": False, - "comment": None, + "name": "d2", + "nullable": True, + "type": VARCHAR(length=20), }, { - "name": "parent_user_id", - "type": INTEGER(), - "nullable": True, - "default": None, "autoincrement": False, + "comment": "Comment\nwith\nescapes", + "default": None, "is_hidden": False, - "comment": None, + "name": "d3", + "nullable": True, + "type": VARCHAR(length=42), }, ], - (None, "comment_test"): [ + (None, "dingalings"): [ { - "name": "id", - "type": INTEGER(), - "nullable": False, - "default": "unique_rowid()", "autoincrement": True, + "comment": None, + "default": "unique_rowid()", "is_hidden": False, - "comment": "id comment", + "name": "dingaling_id", + "nullable": False, + "type": INTEGER(), }, { - "name": "data", - "type": VARCHAR(length=20), - "nullable": True, - "default": None, "autoincrement": False, + "comment": None, + "default": None, "is_hidden": False, - "comment": "data % comment", + "name": "address_id", + "nullable": True, + "type": INTEGER(), }, { - "name": "d2", - "type": VARCHAR(length=20), - "nullable": True, - "default": None, "autoincrement": False, + "comment": None, + "default": None, "is_hidden": False, - "comment": "Comment types type speedily ' \" \\ '' Fun!", + "name": "id_user", + "nullable": True, + "type": INTEGER(), }, { - "name": "d3", - "type": VARCHAR(length=42), + "autoincrement": False, + "comment": None, + "default": None, + "is_hidden": False, + "name": "data", "nullable": True, + "type": VARCHAR(length=30), + }, + ], + (None, "email_addresses"): [ + { + "autoincrement": True, + "comment": None, + "default": "unique_rowid()", + "is_hidden": False, + "name": "address_id", + "nullable": False, + "type": INTEGER(), + }, + { + "autoincrement": False, + "comment": None, "default": None, + "is_hidden": False, + "name": "remote_user_id", + "nullable": True, + "type": INTEGER(), + }, + { "autoincrement": False, + "comment": None, + "default": None, "is_hidden": False, - "comment": "Comment\nwith\rescapes", + "name": "email_address", + "nullable": True, + "type": VARCHAR(length=20), }, ], (None, "no_constraints"): [ { - "name": "data", - "type": VARCHAR(length=20), - "nullable": True, - "default": None, "autoincrement": False, - "is_hidden": False, "comment": None, + "default": None, + "is_hidden": False, + "name": "data", + "nullable": True, + "type": VARCHAR(length=20), } ], (None, "noncol_idx_test_nopk"): [ { - "name": "q", - "type": VARCHAR(length=5), - "nullable": True, - "default": None, "autoincrement": False, - "is_hidden": False, "comment": None, + "default": None, + "is_hidden": False, + "name": "q", + "nullable": True, + "type": VARCHAR(length=5), } ], (None, "noncol_idx_test_pk"): [ { - "name": "id", - "type": INTEGER(), - "nullable": False, - "default": "unique_rowid()", "autoincrement": True, - "is_hidden": False, "comment": None, + "default": "unique_rowid()", + "is_hidden": False, + "name": "id", + "nullable": False, + "type": INTEGER(), }, { - "name": "q", - "type": VARCHAR(length=5), - "nullable": True, - "default": None, "autoincrement": False, - "is_hidden": False, "comment": None, + "default": None, + "is_hidden": False, + "name": "q", + "nullable": True, + "type": VARCHAR(length=5), }, ], - (None, "email_addresses"): [ + (None, "users"): [ { - "name": "address_id", - "type": INTEGER(), - "nullable": False, - "default": "unique_rowid()", "autoincrement": True, - "is_hidden": False, "comment": None, + "default": "unique_rowid()", + "is_hidden": False, + "name": "user_id", + "nullable": False, + "type": INTEGER(), }, { - "name": "remote_user_id", - "type": INTEGER(), - "nullable": True, - "default": None, "autoincrement": False, - "is_hidden": False, "comment": None, - }, - { - "name": "email_address", - "type": VARCHAR(length=20), - "nullable": True, "default": None, - "autoincrement": False, "is_hidden": False, - "comment": None, - }, - ], - (None, "dingalings"): [ - { - "name": "dingaling_id", - "type": INTEGER(), + "name": "test1", "nullable": False, - "default": "unique_rowid()", - "autoincrement": True, - "is_hidden": False, - "comment": None, + "type": CHAR(length=5), }, { - "name": "address_id", - "type": INTEGER(), - "nullable": True, - "default": None, "autoincrement": False, - "is_hidden": False, "comment": None, - }, - { - "name": "id_user", - "type": INTEGER(), - "nullable": True, "default": None, - "autoincrement": False, "is_hidden": False, - "comment": None, + "name": "test2", + "nullable": False, + "type": DOUBLE_PRECISION(precision=53), }, { - "name": "data", - "type": VARCHAR(length=30), - "nullable": True, - "default": None, "autoincrement": False, - "is_hidden": False, "comment": None, + "default": None, + "is_hidden": False, + "name": "parent_user_id", + "nullable": True, + "type": INTEGER(), }, ], } @@ -242,12 +242,12 @@ def test_get_multi_indexes(self): "column_names": ["data"], "column_sorting": {"data": ("nulls_first",)}, "dialect_options": { - "postgresql_ops": { - "data": None, - }, + "postgresql_include": [], + "postgresql_ops": {"data": None}, "postgresql_using": "prefix", }, "duplicates_constraint": "dingalings_data_key", + "include_columns": [], "name": "dingalings_data_key", "unique": True, }, @@ -258,13 +258,12 @@ def test_get_multi_indexes(self): "dingaling_id": ("nulls_first",), }, "dialect_options": { - "postgresql_ops": { - "address_id": None, - "dingaling_id": None, - }, + "postgresql_include": [], + "postgresql_ops": {"address_id": None, "dingaling_id": None}, "postgresql_using": "prefix", }, "duplicates_constraint": "zz_dingalings_multiple", + "include_columns": [], "name": "zz_dingalings_multiple", "unique": True, }, @@ -274,11 +273,11 @@ def test_get_multi_indexes(self): "column_names": ["email_address"], "column_sorting": {"email_address": ("nulls_first",)}, "dialect_options": { - "postgresql_ops": { - "email_address": None, - }, + "postgresql_include": [], + "postgresql_ops": {"email_address": None}, "postgresql_using": "prefix", }, + "include_columns": [], "name": "ix_email_addresses_email_address", "unique": False, } @@ -289,11 +288,11 @@ def test_get_multi_indexes(self): "column_names": ["q"], "column_sorting": {"q": ("desc", "nulls_last")}, "dialect_options": { - "postgresql_ops": { - "q": None, - }, + "postgresql_include": [], + "postgresql_ops": {"q": None}, "postgresql_using": "prefix", }, + "include_columns": [], "name": "noncol_idx_nopk", "unique": False, } @@ -303,11 +302,11 @@ def test_get_multi_indexes(self): "column_names": ["q"], "column_sorting": {"q": ("desc", "nulls_last")}, "dialect_options": { - "postgresql_ops": { - "q": None, - }, + "postgresql_include": [], + "postgresql_ops": {"q": None}, "postgresql_using": "prefix", }, + "include_columns": [], "name": "noncol_idx_pk", "unique": False, } @@ -321,13 +320,11 @@ def test_get_multi_indexes(self): "user_id": ("nulls_first",), }, "dialect_options": { - "postgresql_ops": { - "test1": None, - "test2": None, - "user_id": None, - }, + "postgresql_include": [], + "postgresql_ops": {"test1": None, "test2": None, "user_id": None}, "postgresql_using": "prefix", }, + "include_columns": [], "name": "users_all_idx", "unique": False, }, @@ -335,13 +332,12 @@ def test_get_multi_indexes(self): "column_names": ["test1", "test2"], "column_sorting": {"test1": ("nulls_first",), "test2": ("nulls_first",)}, "dialect_options": { - "postgresql_ops": { - "test1": None, - "test2": None, - }, + "postgresql_include": [], + "postgresql_ops": {"test1": None, "test2": None}, "postgresql_using": "prefix", }, "duplicates_constraint": "users_t_idx", + "include_columns": [], "name": "users_t_idx", "unique": True, }, @@ -358,36 +354,43 @@ def test_get_multi_pk_constraint(self): (None, "comment_test"): { "comment": None, "constrained_columns": ["id"], + "dialect_options": {"postgresql_include": []}, "name": "comment_test_pkey", }, (None, "dingalings"): { "comment": None, "constrained_columns": ["dingaling_id"], + "dialect_options": {"postgresql_include": []}, "name": "dingalings_pkey", }, (None, "email_addresses"): { "comment": "ea pk comment", "constrained_columns": ["address_id"], + "dialect_options": {"postgresql_include": []}, "name": "email_ad_pk", }, (None, "no_constraints"): { "comment": None, "constrained_columns": ["rowid"], + "dialect_options": {"postgresql_include": []}, "name": "no_constraints_pkey", }, (None, "noncol_idx_test_nopk"): { "comment": None, "constrained_columns": ["rowid"], + "dialect_options": {"postgresql_include": []}, "name": "noncol_idx_test_nopk_pkey", }, (None, "noncol_idx_test_pk"): { "comment": None, "constrained_columns": ["id"], + "dialect_options": {"postgresql_include": []}, "name": "noncol_idx_test_pk_pkey", }, (None, "users"): { "comment": None, "constrained_columns": ["user_id"], + "dialect_options": {"postgresql_include": []}, "name": "users_pkey", }, }, From 094c7047ae17181ad5767208fb5b157dd57c7bba Mon Sep 17 00:00:00 2001 From: Gord Thompson Date: Thu, 23 Apr 2026 13:53:24 -0600 Subject: [PATCH 2/4] Update CHANGES.md and requirements --- CHANGES.md | 2 +- dev-requirements.txt | 16 ++++++++-------- test-requirements.txt | 4 ++-- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 8ba8792..b659fe6 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -1,6 +1,7 @@ # Version 2.0.5 Unreleased +- Fix get_columns() to include identity column info (#297) # Version 2.0.4 April 23, 2026 @@ -9,7 +10,6 @@ April 23, 2026 - Fix reflection of TIMESTAMPTZ columns (#276), thanks to @nvachhar - Fix reflection of JSONB columns (#277) - Fix compatibility issues with Alembic 1.18 (via SQLA 2.0.47) -- Fix get_columns() to include identity column info (#297) - Update minimum Python version to 3.10 - Compile MySQL-style `func.timestampdiff(unit, start, end)` to a PostgreSQL-style `EXTRACT(EPOCH FROM ...)` expression on the cockroachdb diff --git a/dev-requirements.txt b/dev-requirements.txt index aa31d0f..c3a0f58 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -1,12 +1,12 @@ backports-tarfile==1.2.0 # via jaraco-context -cachetools==7.0.5 +cachetools==7.0.6 # via tox -certifi==2026.2.25 +certifi==2026.4.22 # via requests cffi==2.0.0 # via cryptography -chardet==7.4.1 +chardet==7.4.3 # via tox charset-normalizer==3.4.7 # via requests @@ -18,14 +18,14 @@ distlib==0.4.0 # via virtualenv docutils==0.22.4 # via readme-renderer -filelock==3.25.2 +filelock==3.29.0 # via # python-discovery # tox # virtualenv id==1.6.1 # via twine -idna==3.11 +idna==3.13 # via requests importlib-metadata==9.0.0 # via keyring @@ -51,7 +51,7 @@ more-itertools==11.0.2 # jaraco-functools nh3==0.3.4 # via readme-renderer -packaging==26.0 +packaging==26.1 # via # pyproject-api # tox @@ -105,7 +105,7 @@ urllib3==2.6.3 # id # requests # twine -virtualenv==21.2.1 +virtualenv==21.2.4 # via tox -zipp==3.23.0 +zipp==3.23.1 # via importlib-metadata diff --git a/test-requirements.txt b/test-requirements.txt index 8f5c688..cae53c3 100644 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -20,13 +20,13 @@ mock==5.2.0 # via -r test-requirements.in more-itertools==11.0.2 # via -r test-requirements.in -packaging==26.0 +packaging==26.1 # via pytest pluggy==1.6.0 # via pytest psycopg==3.3.3 # via -r test-requirements.in -psycopg2==2.9.11 +psycopg2==2.9.12 # via -r test-requirements.in pygments==2.20.0 # via pytest From 6f02b86540dc98090409e205c3ae2aa948f545f6 Mon Sep 17 00:00:00 2001 From: Gord Thompson Date: Sat, 16 May 2026 08:05:15 -0600 Subject: [PATCH 3/4] Apply changes from review --- CHANGES.md | 4 +- test/test_column_reflect.py | 76 +++++++++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+), 1 deletion(-) diff --git a/CHANGES.md b/CHANGES.md index b659fe6..6d48e70 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -1,7 +1,9 @@ # Version 2.0.5 Unreleased -- Fix get_columns() to include identity column info (#297) +- Re-work get_multi_columns() to + - include identity column info (#297) + - avoid parse error when reflecting ENUMs (#303) # Version 2.0.4 April 23, 2026 diff --git a/test/test_column_reflect.py b/test/test_column_reflect.py index 33df0cf..090c274 100644 --- a/test/test_column_reflect.py +++ b/test/test_column_reflect.py @@ -8,6 +8,7 @@ inspect, BigInteger, Identity, + Computed, ) from sqlalchemy.testing import fixtures, eq_ @@ -33,6 +34,23 @@ Column("txt", String), ) +with_identity_always = Table( + "with_identity_always", + meta, + Column("id", BigInteger, Identity(always=True), primary_key=True), + Column("txt", String), +) + +with_computed_stored = Table( + "with_computed_stored", + meta, + Column("id", Integer, primary_key=True), + Column("id2", Integer, Computed("id + 1", persisted=True)), +) + +# Note: CockroachDB computed columns do not support 'virtual' persistence; +# set the 'persisted' flag to None or True for CockroachDB support. + class ReflectHiddenColumnsTest(fixtures.TestBase): __requires__ = ("sync_driver",) @@ -147,3 +165,61 @@ def test_reflect_identity(self): }, ], ) + eq_( + self._get_col_info("with_identity_always"), + [ + { + "autoincrement": True, + "comment": None, + "default": "nextval('public.with_identity_always_id_seq'::REGCLASS)", + "identity": { + "always": True, + "cache": 1, + "cycle": False, + "increment": 1, + "maxvalue": 9223372036854775807, + "minvalue": 1, + "start": 1, + }, + "is_hidden": False, + "name": "id", + "nullable": False, + "type": "INTEGER", + }, + { + "autoincrement": False, + "comment": None, + "default": None, + "is_hidden": False, + "name": "txt", + "nullable": True, + "type": "VARCHAR", + }, + ], + ) + + def test_reflect_computed_stored(self): + eq_( + self._get_col_info("with_computed_stored"), + [ + { + "autoincrement": True, + "comment": None, + "default": "unique_rowid()", + "is_hidden": False, + "name": "id", + "nullable": False, + "type": "INTEGER", + }, + { + "autoincrement": False, + "comment": None, + "computed": {"persisted": True, "sqltext": "id + 1"}, + "default": None, + "is_hidden": False, + "name": "id2", + "nullable": True, + "type": "INTEGER", + }, + ], + ) From 50cfcf17c3626a0876e2abd2f181f2125628ac13 Mon Sep 17 00:00:00 2001 From: Gord Thompson Date: Sat, 16 May 2026 09:39:44 -0600 Subject: [PATCH 4/4] Apply changes from review (cont'd - 2) --- sqlalchemy_cockroachdb/base.py | 67 ++++------------------------------ test/test_column_reflect.py | 36 ++++++++++++++++++ 2 files changed, 43 insertions(+), 60 deletions(-) diff --git a/sqlalchemy_cockroachdb/base.py b/sqlalchemy_cockroachdb/base.py index 927fa32..d595305 100644 --- a/sqlalchemy_cockroachdb/base.py +++ b/sqlalchemy_cockroachdb/base.py @@ -15,64 +15,11 @@ from sqlalchemy import VARCHAR from sqlalchemy.dialects.postgresql.base import PGDialect from sqlalchemy.dialects.postgresql import BYTEA -from sqlalchemy.dialects.postgresql import INET -from sqlalchemy.dialects.postgresql import JSONB -from sqlalchemy.dialects.postgresql import UUID from sqlalchemy.ext.compiler import compiles -import sqlalchemy.types as sqltypes - from .stmt_compiler import CockroachCompiler, CockroachIdentifierPreparer from .ddl_compiler import CockroachDDLCompiler -# Map type names (as returned by information_schema) to sqlalchemy type -# objects. -# -# TODO(bdarnell): test more of these. The stock test suite only covers -# a few basic ones. -_type_map = { - "bool": sqltypes.BOOLEAN, # introspection returns "BOOL" not boolean - "boolean": sqltypes.BOOLEAN, - "bigint": sqltypes.INT, - "int": sqltypes.INT, - "int2": sqltypes.INT, - "int4": sqltypes.INT, - "int64": sqltypes.INT, - "int8": sqltypes.INT, - "integer": sqltypes.INT, - "smallint": sqltypes.INT, - "double precision": sqltypes.FLOAT, - "float": sqltypes.FLOAT, - "float4": sqltypes.FLOAT, - "float8": sqltypes.FLOAT, - "real": sqltypes.FLOAT, - "dec": sqltypes.DECIMAL, - "decimal": sqltypes.DECIMAL, - "numeric": sqltypes.DECIMAL, - "date": sqltypes.DATE, - "time": sqltypes.Time, - "time without time zone": sqltypes.Time, - "timestamp": sqltypes.TIMESTAMP, - "timestamptz": sqltypes.TIMESTAMP(timezone=True), - "timestamp with time zone": sqltypes.TIMESTAMP(timezone=True), - "timestamp without time zone": sqltypes.TIMESTAMP, - "interval": sqltypes.Interval, - "char": sqltypes.CHAR, - "char varying": sqltypes.VARCHAR, - "character": sqltypes.CHAR, - "character varying": sqltypes.VARCHAR, - "string": sqltypes.VARCHAR, - "text": sqltypes.VARCHAR, - "varchar": sqltypes.VARCHAR, - "blob": sqltypes.BLOB, - "bytea": sqltypes.BLOB, - "bytes": sqltypes.BLOB, - "json": sqltypes.JSON, - "jsonb": JSONB, - "uuid": UUID, - "inet": INET, -} - class _SavepointState(threading.local): """Hack to override names used in savepoint statements. @@ -213,7 +160,7 @@ def get_multi_columns(self, connection, schema, filter_names, scope, kind, **kw) .all() ) is_hidden = {x["column_name"]: x["is_hidden"] for x in table_columns} - for col in columns: + for col in columns[:]: if is_hidden[col["name"]] and not _include_hidden: columns.remove(col) else: @@ -221,19 +168,19 @@ def get_multi_columns(self, connection, schema, filter_names, scope, kind, **kw) if col["default"] == "unique_rowid()": col["autoincrement"] = True if isinstance(col["type"], BIGINT): - col["type"] = INTEGER + col["type"] = INTEGER() elif isinstance(col["type"], ARRAY) and isinstance( col["type"].item_type, BIGINT ): col["type"].item_type = INTEGER() elif isinstance(col["type"], BYTEA): - col["type"] = BLOB + col["type"] = BLOB() elif isinstance(col["type"], ARRAY) and isinstance( col["type"].item_type, BYTEA ): col["type"].item_type = BLOB() elif isinstance(col["type"], DOUBLE_PRECISION): - col["type"] = FLOAT + col["type"] = FLOAT() elif isinstance(col["type"], ARRAY) and isinstance( col["type"].item_type, DOUBLE_PRECISION ): @@ -247,19 +194,19 @@ def get_multi_columns(self, connection, schema, filter_names, scope, kind, **kw) col["type"].item_type.precision, col["type"].item_type.scale ) elif isinstance(col["type"], REAL): - col["type"] = FLOAT + col["type"] = FLOAT() elif isinstance(col["type"], ARRAY) and isinstance( col["type"].item_type, REAL ): col["type"].item_type = FLOAT() elif isinstance(col["type"], SMALLINT): - col["type"] = INTEGER + col["type"] = INTEGER() elif isinstance(col["type"], ARRAY) and isinstance( col["type"].item_type, SMALLINT ): col["type"].item_type = INTEGER() elif isinstance(col["type"], TEXT): - col["type"] = VARCHAR + col["type"] = VARCHAR() elif isinstance(col["type"], ARRAY) and isinstance( col["type"].item_type, TEXT ): diff --git a/test/test_column_reflect.py b/test/test_column_reflect.py index 090c274..8431efd 100644 --- a/test/test_column_reflect.py +++ b/test/test_column_reflect.py @@ -21,6 +21,14 @@ Column("txt", String), ) +with_pk_other_schema = Table( + "with_pk_other_schema", + meta, + Column("id", Integer, primary_key=True), + Column("txt", String), + schema="test_schema", +) + without_pk = Table( "without_pk", meta, @@ -165,6 +173,7 @@ def test_reflect_identity(self): }, ], ) + eq_( self._get_col_info("with_identity_always"), [ @@ -223,3 +232,30 @@ def test_reflect_computed_stored(self): }, ], ) + + # def test_reflect_other_schema(self): + # # TODO: uncomment when resolved + # # verify https://github.com/cockroachdb/cockroach/issues/170049 + # eq_( + # self._get_col_info("with_pk_other_schema"), + # [ + # { + # "autoincrement": True, + # "comment": None, + # "default": "unique_rowid()", + # "is_hidden": False, + # "name": "id", + # "nullable": False, + # "type": "INTEGER", + # }, + # { + # "autoincrement": False, + # "comment": None, + # "default": None, + # "is_hidden": False, + # "name": "txt", + # "nullable": True, + # "type": "VARCHAR", + # }, + # ], + # )