index with data support

caretdev · Nov 15, 2022 · 4a1f299 · 4a1f299
1 parent 389dcaa
commit 4a1f299
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 2 deletions.
diff --git a/sqlalchemy_iris/base.py b/sqlalchemy_iris/base.py
@@ -579,6 +579,30 @@ def get_column_specification(self, column, **kwargs):
     def post_create_table(self, table):
         return " WITH %CLASSPARAMETER ALLOWIDENTITYINSERT = 1"
 
+    def visit_create_index(
+        self, create, include_schema=False, include_table_schema=True, **kw
+    ):
+        text = super().visit_create_index(create, include_schema, include_table_schema, **kw)
+
+        index = create.element
+        preparer = self.preparer
+
+        # handle other included columns
+        includeclause = index.dialect_options["iris"]["include"]
+        if includeclause:
+            inclusions = [
+                index.table.c[col]
+                if isinstance(col, util.string_types)
+                else col
+                for col in includeclause
+            ]
+
+            text += " WITH DATA (%s)" % ", ".join(
+                [preparer.quote(c.name) for c in inclusions]
+            )
+
+        return text
+
 
 class IRISTypeCompiler(compiler.GenericTypeCompiler):
     def visit_boolean(self, type_, **kw):
@@ -667,6 +691,10 @@ class IRISDialect(default.DefaultDialect):
     type_compiler = IRISTypeCompiler
     execution_ctx_cls = IRISExecutionContext
 
+    construct_arguments = [
+        (schema.Index, {"include": None}),
+    ]
+
     def __init__(self, **kwargs):
         default.DefaultDialect.__init__(self, **kwargs)
 
@@ -808,6 +836,8 @@ def has_table(self, connection, table_name, schema=None, **kw):
     def get_indexes(self, connection, table_name, schema=None, unique=False, **kw):
         schema_name = self.get_schema(schema)
         indexes = ischema.indexes
+        tables = ischema.tables
+        index_def = ischema.index_definition
 
         s = (
             sql.select(
@@ -816,6 +846,20 @@ def get_indexes(self, connection, table_name, schema=None, unique=False, **kw):
                 indexes.c.primary_key,
                 indexes.c.non_unique,
                 indexes.c.asc_or_desc,
+                index_def.c.Data,
+            )
+            .select_from(indexes)
+            .outerjoin(
+                index_def,
+                sql.and_(
+                    index_def.c.SqlName == indexes.c.index_name,
+                    index_def.c.parent ==
+                    sql.select(tables.c.classname)
+                    .where(
+                        indexes.c.table_name == tables.c.table_name,
+                        indexes.c.table_schema == tables.c.table_schema,
+                    ).scalar_subquery()
+                ),
             )
             .where(
                 sql.and_(
@@ -838,6 +882,7 @@ def get_indexes(self, connection, table_name, schema=None, unique=False, **kw):
                 _,
                 nuniq,
                 _,
+                include,
             ) = row
 
             indexrec = indexes[idxname]
@@ -849,6 +894,12 @@ def get_indexes(self, connection, table_name, schema=None, unique=False, **kw):
             indexrec["column_names"].append(
                 self.normalize_name(colname)
             )
+            include = include.split(',') if include else []
+            indexrec["include_columns"] = include
+            if include:
+                indexrec["dialect_options"] = {
+                    "iris_include": include
+                }
 
         indexes = list(indexes.values())
         return indexes
@@ -1007,7 +1058,7 @@ def get_columns(self, connection, table_name, schema=None, **kw):
         schema_name = self.get_schema(schema)
         tables = ischema.tables
         columns = ischema.columns
-        property = ischema.property
+        property = ischema.property_definition
 
         whereclause = sql.and_(
             columns.c.table_name == str(table_name),

diff --git a/sqlalchemy_iris/information_schema.py b/sqlalchemy_iris/information_schema.py
@@ -72,7 +72,7 @@ def process_result_value(self, value, dialect):
     Column("DESCIPTION", String, key="desciption"),
     schema="INFORMATION_SCHEMA",
 )
-property = Table(
+property_definition = Table(
     "PropertyDefinition",
     ischema,
     Column("parent", String),
@@ -101,6 +101,15 @@ def process_result_value(self, value, dialect):
     schema="INFORMATION_SCHEMA",
 )
 
+index_definition = Table(
+    "IndexDefinition",
+    ischema,
+    Column("parent", String),
+    Column("SqlName", String),
+    Column("Data", String),
+    schema="%Dictionary",
+)
+
 key_constraints = Table(
     "KEY_COLUMN_USAGE",
     ischema,