Update DESCRIPTION.md and add support for indexes

.github/workflows/build_test.yml

@@ -111,6 +111,9 @@ jobs:
         run: |
           gpg --quiet --batch --yes --decrypt --passphrase="$PARAMETERS_SECRET" \
           .github/workflows/parameters/parameters_${{ matrix.cloud-provider }}.py.gpg > tests/parameters.py
+      - name: Run test for AWS
+        run: hatch run test-dialect-aws
+        if: matrix.cloud-provider == 'aws'
       - name: Run tests
         run: hatch run test-dialect
       - uses: actions/upload-artifact@v4
@@ -203,6 +206,9 @@ jobs:
           python -m pip install -U uv
           python -m uv pip install -U hatch
           python -m hatch env create default
+      - name: Run test for AWS
+        run: hatch run sa14:test-dialect-aws
+        if: matrix.cloud-provider == 'aws'
       - name: Run tests
         run: hatch run sa14:test-dialect
       - uses: actions/upload-artifact@v4

DESCRIPTION.md

@@ -12,7 +12,7 @@ Source code is also available at:
 - (Unreleased)
 
   - Add support for dynamic tables and required options
-  - Fixed SAWarning when registering functions with existing name in default namespace
+  - Add support for hybrid tables
   - Fixed SAWarning when registering functions with existing name in default namespace
 
 - v1.6.1(July 9, 2024)

pyproject.toml

@@ -92,6 +92,7 @@ SQLACHEMY_WARN_20 = "1"
 check = "pre-commit run --all-files"
 test-dialect = "pytest -ra -vvv --tb=short --cov snowflake.sqlalchemy --cov-append --junitxml ./junit.xml --ignore=tests/sqlalchemy_test_suite tests/"
 test-dialect-compatibility = "pytest -ra -vvv --tb=short --cov snowflake.sqlalchemy --cov-append --junitxml ./junit.xml tests/sqlalchemy_test_suite"
+test-dialect-aws = "pytest -m \"aws\" -ra -vvv --tb=short --cov snowflake.sqlalchemy --cov-append --junitxml ./junit.xml --ignore=tests/sqlalchemy_test_suite tests/"
 gh-cache-sum = "python -VV | sha256sum | cut -d' ' -f1"
 check-import = "python -c 'import snowflake.sqlalchemy; print(snowflake.sqlalchemy.__version__)'"
 
@@ -110,7 +111,7 @@ line-length = 88
 line-length = 88
 
 [tool.pytest.ini_options]
-addopts = "-m 'not feature_max_lob_size'"
+addopts = "-m 'not feature_max_lob_size and not aws'"
 markers = [
   # Optional dependency groups markers
   "lambda: AWS lambda tests",

src/snowflake/sqlalchemy/snowdialect.py

@@ -65,6 +65,7 @@
     _CUSTOM_Float,
     _CUSTOM_Time,
 )
+from .sql.custom_schema.custom_table_prefix import CustomTablePrefix
 from .util import (
     _update_connection_application_name,
     parse_url_boolean,
@@ -898,15 +899,21 @@ def get_multi_indexes(
         """
         Gets the indexes definition
         """
+
+        table_prefixes = self.get_multi_prefixes(
+            connection, schema, filter_prefix=CustomTablePrefix.HYBRID.name
+        )
+        if len(table_prefixes) == 0:
+            return []
         schema = schema or self.default_schema_name
         if not schema:
             result = connection.execute(
-                text("SHOW /* sqlalchemy:get_view_definition */ INDEXES")
+                text("SHOW /* sqlalchemy:get_multi_indexes */ INDEXES")
             )
         else:
             result = connection.execute(
                 text(
-                    f"SHOW /* sqlalchemy:get_indexes */ INDEXES IN SCHEMA {self._denormalize_quote_join(schema)}"
+                    f"SHOW /* sqlalchemy:get_multi_indexes */ INDEXES IN SCHEMA {self._denormalize_quote_join(schema)}"
                 )
             )
 
@@ -918,6 +925,12 @@ def get_multi_indexes(
             if (
                 row[n2i["name"]] == f'SYS_INDEX_{row[n2i["table"]]}_PRIMARY'
                 or table not in filter_names
+                or (schema, table) not in table_prefixes
+                or (
+                    (schema, table) in table_prefixes
+                    and CustomTablePrefix.HYBRID.name
+                    not in table_prefixes[(schema, table)]
+                )
             ):
                 continue
             index = {
@@ -942,6 +955,50 @@ def _value_or_default(self, data, table, schema):
         else:
             return []
 
+    def get_prefixes_from_data(self, n2i, row, **kw):
+        prefixes_found = []
+        for valid_prefix in CustomTablePrefix:
+            key = f"is_{valid_prefix.name.lower()}"
+            if key in n2i and row[n2i[key]] == "Y":
+                prefixes_found.append(valid_prefix.name)
+        return prefixes_found
+
+    @reflection.cache
+    def get_multi_prefixes(
+        self, connection, schema, table_name=None, filter_prefix=None, **kw
+    ):
+        """
+        Gets all table prefixes
+        """
+        schema = schema or self.default_schema_name
+        filter = f"LIKE '{table_name}'" if table_name else ""
+        if schema:
+            result = connection.execute(
+                text(
+                    f"SHOW /* sqlalchemy:get_multi_prefixes */ {filter} TABLES IN SCHEMA {schema}"
+                )
+            )
+        else:
+            result = connection.execute(
+                text(
+                    f"SHOW /* sqlalchemy:get_multi_prefixes */ {filter} TABLES LIKE '{table_name}'"
+                )
+            )
+
+        n2i = self.__class__._map_name_to_idx(result)
+        tables_prefixes = {}
+        for row in result.cursor.fetchall():
+            table = self.normalize_name(str(row[n2i["name"]]))
+            table_prefixes = self.get_prefixes_from_data(n2i, row)
+            if filter_prefix and filter_prefix not in table_prefixes:
+                continue
+            if (schema, table) in tables_prefixes:
+                tables_prefixes[(schema, table)].append(table_prefixes)
+            else:
+                tables_prefixes[(schema, table)] = table_prefixes
+
+        return tables_prefixes
+
     @reflection.cache
     def get_indexes(self, connection, tablename, schema, **kw):
         """
@@ -976,7 +1033,7 @@ def check_table(table, connection, _ddl_runner, **kw):
     if HybridTable.is_equal_type(table):  # noqa
         return True
     if isinstance(_ddl_runner.dialect, SnowflakeDialect) and table.indexes:
-        raise NotImplementedError("Snowflake does not support indexes")
+        raise NotImplementedError("Only Snowflake Hybrid Tables supports indexes")
 
 
 dialect = SnowflakeDialect

src/snowflake/sqlalchemy/sql/custom_schema/custom_table_base.py

@@ -10,12 +10,17 @@
 from ..._constants import DIALECT_NAME
 from ...compat import IS_VERSION_20
 from ...custom_commands import NoneType
+from .custom_table_prefix import CustomTablePrefix
 from .options.table_option import TableOption
 
 
 class CustomTableBase(Table):
-    __table_prefix__ = ""
-    _support_primary_and_foreign_keys = True
+    __table_prefixes__: typing.List[CustomTablePrefix] = []
+    _support_primary_and_foreign_keys: bool = True
+
+    @property
+    def table_prefixes(self) -> typing.List[str]:
+        return [prefix.name for prefix in self.__table_prefixes__]
 
     def __init__(
         self,
@@ -24,8 +29,8 @@ def __init__(
         *args: SchemaItem,
         **kw: Any,
     ) -> None:
-        if self.__table_prefix__ != "":
-            prefixes = kw.get("prefixes", []) + [self.__table_prefix__]
+        if len(self.__table_prefixes__) > 0:
+            prefixes = kw.get("prefixes", []) + self.table_prefixes
             kw.update(prefixes=prefixes)
         if not IS_VERSION_20 and hasattr(super(), "_init"):
             super()._init(name, metadata, *args, **kw)
@@ -40,7 +45,7 @@ def _validate_table(self):
             self.primary_key or self.foreign_keys
         ):
             raise ArgumentError(
-                f"Primary key and foreign keys are not supported in {self.__table_prefix__} TABLE."
+                f"Primary key and foreign keys are not supported in {' '.join(self.table_prefixes)} TABLE."
             )
 
         return True
@@ -52,9 +57,8 @@ def _get_dialect_option(self, option_name: str) -> typing.Optional[TableOption]:
 
     @classmethod
     def is_equal_type(cls, table: Table) -> bool:
-        if isinstance(table, cls.__class__):
-            return True
-        for prefix in table._prefixes:
-            if prefix == cls.__table_prefix__:
-                return True
-        return False
+        for prefix in cls.__table_prefixes__:
+            if prefix.name not in table._prefixes:
+                return False
+
+        return True

src/snowflake/sqlalchemy/sql/custom_schema/custom_table_prefix.py

@@ -0,0 +1,13 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+
+from enum import Enum
+
+
+class CustomTablePrefix(Enum):
+    DEFAULT = 0
+    EXTERNAL = 1
+    EVENT = 2
+    HYBRID = 3
+    ICEBERG = 4
+    DYNAMIC = 5

src/snowflake/sqlalchemy/sql/custom_schema/dynamic_table.py

@@ -10,6 +10,7 @@
 
 from snowflake.sqlalchemy.custom_commands import NoneType
 
+from .custom_table_prefix import CustomTablePrefix
 from .options.target_lag import TargetLag
 from .options.warehouse import Warehouse
 from .table_from_query import TableFromQueryBase
@@ -27,7 +28,7 @@ class DynamicTable(TableFromQueryBase):
 
     """
 
-    __table_prefix__ = "DYNAMIC"
+    __table_prefixes__ = [CustomTablePrefix.DYNAMIC]
 
     _support_primary_and_foreign_keys = False
 

src/snowflake/sqlalchemy/sql/custom_schema/hybrid_table.py

@@ -10,6 +10,7 @@
 from snowflake.sqlalchemy.custom_commands import NoneType
 
 from .custom_table_base import CustomTableBase
+from .custom_table_prefix import CustomTablePrefix
 
 
 class HybridTable(CustomTableBase):
@@ -22,7 +23,7 @@ class HybridTable(CustomTableBase):
     interface for creating dynamic tables and management.
     """
 
-    __table_prefix__ = "HYBRID"
+    __table_prefixes__ = [CustomTablePrefix.HYBRID]
 
     _support_primary_and_foreign_keys = True
 

tests/__snapshots__/test_orm.ambr

@@ -0,0 +1,4 @@
+# serializer version: 1
+# name: test_orm_one_to_many_relationship_with_hybrid_table
+  ProgrammingError('(snowflake.connector.errors.ProgrammingError) 200009 (22000): Foreign key constraint "SYS_INDEX_HB_TBL_ADDRESS_FOREIGN_KEY_USER_ID_HB_TBL_USER_ID" was violated.')
+# ---

tests/custom_tables/__snapshots__/test_compile_dynamic_table.ambr

@@ -0,0 +1,13 @@
+# serializer version: 1
+# name: test_compile_dynamic_table
+  "CREATE DYNAMIC TABLE test_dynamic_table (\tid INTEGER, \tgeom GEOMETRY)\tWAREHOUSE = warehouse\tTARGET_LAG = '10 seconds'\tAS SELECT * FROM table"
+# ---
+# name: test_compile_dynamic_table_orm
+  "CREATE DYNAMIC TABLE test_dynamic_table_orm (\tid INTEGER, \tname VARCHAR)\tWAREHOUSE = warehouse\tTARGET_LAG = '10 seconds'\tAS SELECT * FROM table"
+# ---
+# name: test_compile_dynamic_table_orm_with_str_keys
+  "CREATE DYNAMIC TABLE test_dynamic_table_orm_2 (\tid INTEGER, \tname VARCHAR)\tWAREHOUSE = warehouse\tTARGET_LAG = '10 seconds'\tAS SELECT * FROM table"
+# ---
+# name: test_compile_dynamic_table_with_selectable
+  "CREATE DYNAMIC TABLE dynamic_test_table_1 (\tid INTEGER, \tname VARCHAR)\tWAREHOUSE = warehouse\tTARGET_LAG = '10 seconds'\tAS SELECT test_table_1.id, test_table_1.name FROM test_table_1 WHERE test_table_1.id = 23"
+# ---

tests/custom_tables/__snapshots__/test_compile_hybrid_table.ambr

@@ -0,0 +1,7 @@
+# serializer version: 1
+# name: test_compile_hybrid_table
+  'CREATE HYBRID TABLE test_hybrid_table (\tid INTEGER NOT NULL AUTOINCREMENT, \tname VARCHAR, \tgeom GEOMETRY, \tPRIMARY KEY (id))'
+# ---
+# name: test_compile_hybrid_table_orm
+  'CREATE HYBRID TABLE test_hybrid_table_orm (\tid INTEGER NOT NULL AUTOINCREMENT, \tname VARCHAR, \tPRIMARY KEY (id))'
+# ---

tests/custom_tables/__snapshots__/test_create_hybrid_table.ambr

@@ -0,0 +1,7 @@
+# serializer version: 1
+# name: test_create_hybrid_table
+  "[(1, 'test')]"
+# ---
+# name: test_create_hybrid_table_with_multiple_index
+  ProgrammingError("(snowflake.connector.errors.ProgrammingError) 391480 (0A000): Another index is being built on table 'TEST_HYBRID_TABLE_WITH_MULTIPLE_INDEX'. Only one index can be built at a time. Either cancel the other index creation or wait until it is complete.")
+# ---

tests/custom_tables/__snapshots__/test_reflect_hybrid_table.ambr

@@ -0,0 +1,4 @@
+# serializer version: 1
+# name: test_simple_reflection_hybrid_table_as_table
+  'CREATE TABLE test_hybrid_table_reflection (\tid DECIMAL(38, 0) NOT NULL, \tname VARCHAR(16777216), \tCONSTRAINT demo_name PRIMARY KEY (id))'
+# ---

tests/custom_tables/test_compile_hybrid_table.py

@@ -1,14 +1,15 @@
 #
 # Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
 #
-
+import pytest
 from sqlalchemy import Column, Integer, MetaData, String
 from sqlalchemy.orm import declarative_base
 from sqlalchemy.sql.ddl import CreateTable
 
 from snowflake.sqlalchemy import GEOMETRY, HybridTable
 
 
+@pytest.mark.aws
 def test_compile_hybrid_table(sql_compiler, snapshot):
     metadata = MetaData()
     table_name = "test_hybrid_table"
@@ -27,6 +28,7 @@ def test_compile_hybrid_table(sql_compiler, snapshot):
     assert actual == snapshot
 
 
+@pytest.mark.aws
 def test_compile_hybrid_table_orm(sql_compiler, snapshot):
     Base = declarative_base()
 

tests/custom_tables/test_create_dynamic_table.py

@@ -76,7 +76,6 @@ def test_create_dynamic_table_without_dynamictable_class(
         TargetLag(1, TimeUnit.HOURS),
         Warehouse(warehouse),
         AsQuery("SELECT id, name from test_table_1;"),
-        AsQuery("SELECT id, name from test_table_1;"),
         prefixes=["DYNAMIC"],
     )
 

tests/custom_tables/test_create_hybrid_table.py

@@ -6,9 +6,10 @@
 from sqlalchemy import Column, Index, Integer, MetaData, String, select
 from sqlalchemy.orm import Session, declarative_base
 
-from src.snowflake.sqlalchemy import HybridTable
+from snowflake.sqlalchemy import HybridTable
 
 
+@pytest.mark.aws
 def test_create_hybrid_table(engine_testaccount, db_parameters, snapshot):
     metadata = MetaData()
     table_name = "test_create_hybrid_table"
@@ -24,7 +25,6 @@ def test_create_hybrid_table(engine_testaccount, db_parameters, snapshot):
 
     with engine_testaccount.connect() as conn:
         ins = dynamic_test_table_1.insert().values(id=1, name="test")
-
         conn.execute(ins)
         conn.commit()
 
@@ -37,6 +37,7 @@ def test_create_hybrid_table(engine_testaccount, db_parameters, snapshot):
         metadata.drop_all(engine_testaccount)
 
 
+@pytest.mark.aws
 def test_create_hybrid_table_with_multiple_index(
     engine_testaccount, db_parameters, snapshot, sql_compiler
 ):
@@ -64,6 +65,7 @@ def test_create_hybrid_table_with_multiple_index(
         metadata.drop_all(engine_testaccount)
 
 
+@pytest.mark.aws
 def test_create_hybrid_table_with_orm(sql_compiler, engine_testaccount):
     Base = declarative_base()
     session = Session(bind=engine_testaccount)

tests/custom_tables/test_reflect_hybrid_table.py

@@ -1,10 +1,12 @@
 #
 # Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
 #
+import pytest
 from sqlalchemy import MetaData, Table
 from sqlalchemy.sql.ddl import CreateTable
 
 
+@pytest.mark.aws
 def test_simple_reflection_hybrid_table_as_table(
     engine_testaccount, db_parameters, sql_compiler, snapshot
 ):
@@ -37,6 +39,7 @@ def test_simple_reflection_hybrid_table_as_table(
         metadata.drop_all(engine_testaccount)
 
 
+@pytest.mark.aws
 def test_reflect_hybrid_table_with_index(
     engine_testaccount, db_parameters, sql_compiler
 ):

tests/test_core.py

@@ -506,7 +506,7 @@ def test_get_indexes(engine_testaccount, db_parameters):
     """
     Tests get indexes
 
-    NOTE: Snowflake Tables doesn't support indexes
+    NOTE: Only Snowflake Hybrid Tables support indexes
     """
     schema = db_parameters["schema"]
     metadata = MetaData()