Postgres: ability to create indexes

CHANGELOG.md

@@ -22,6 +22,7 @@
 - Support disabling schema tests, and configuring tests from `dbt_project.yml` ([#3252](https:/fishtown-analytics/dbt/issues/3252),
 [#3253](https:/fishtown-analytics/dbt/issues/3253), [#3257](https:/fishtown-analytics/dbt/pull/3257))
 - Add Jinja tag for tests ([#1173](https:/fishtown-analytics/dbt/issues/1173), [#3261](https:/fishtown-analytics/dbt/pull/3261))
+- Add native support for Postgres index creation ([#804](https:/fishtown-analytics/dbt/issues/804), [3106](https:/fishtown-analytics/dbt/pull/3106))
 
 ### Under the hood
 - Add dependabot configuration for alerting maintainers about keeping dependencies up to date and secure. ([#3061](https:/fishtown-analytics/dbt/issues/3061), [#3062](https:/fishtown-analytics/dbt/pull/3062))
@@ -49,6 +50,7 @@ Contributors:
 - [@cgopalan](https:/cgopalan) ([#3165](https:/fishtown-analytics/dbt/pull/3165), [#3182](https:/fishtown-analytics/dbt/pull/3182))
 - [@fux](https:/fuchsst) ([#3241](https:/fishtown-analytics/dbt/issues/3241))
 - [@dmateusp](https:/dmateusp) ([#3270](https:/fishtown-analytics/dbt/pull/3270))
+- [@arzavj](https:/arzavj) ([3106](https:/fishtown-analytics/dbt/pull/3106))
 
 ## dbt 0.19.1 (March 31, 2021)
 

core/dbt/include/global_project/macros/adapters/common.sql

@@ -51,6 +51,29 @@
 
 {% endmacro %}
 
+{% macro get_create_index_sql(relation, index_dict) -%}
+ {{ return(adapter.dispatch('get_create_index_sql')(relation, index_dict)) }}
+{% endmacro %}
+
+{% macro default__get_create_index_sql(relation, index_dict) -%}
+ {% do return(None) %}
+{% endmacro %}
+
+{% macro create_indexes(relation) -%}
+ {{ adapter.dispatch('create_indexes')(relation) }}
+{%- endmacro %}
+
+{% macro default__create_indexes(relation) -%}
+ {%- set _indexes = config.get('indexes', default=[]) -%}
+
+ {% for _index_dict in _indexes %}
+ {% set create_index_sql = get_create_index_sql(relation, _index_dict) %}
+ {% if create_index_sql %}
+ {% do run_query(create_index_sql) %}
+ {% endif %}
+ {% endfor %}
+{% endmacro %}
+
 {% macro create_view_as(relation, sql) -%}
  {{ adapter.dispatch('create_view_as')(relation, sql) }}
 {%- endmacro %}

core/dbt/include/global_project/macros/materializations/incremental/incremental.sql

@@ -39,6 +39,10 @@
 
  {% do persist_docs(target_relation, model) %}
 
+ {% if existing_relation is none or existing_relation.is_view or should_full_refresh() %}
+ {% do create_indexes(target_relation) %}
+ {% endif %}
+
  {{ run_hooks(post_hooks, inside_transaction=True) }}
 
  -- `COMMIT` happens here

core/dbt/include/global_project/macros/materializations/seed/seed.sql

@@ -142,6 +142,10 @@
  {% set target_relation = this.incorporate(type='table') %}
  {% do persist_docs(target_relation, model) %}
 
+ {% if full_refresh_mode or not exists_as_table %}
+ {% do create_indexes(target_relation) %}
+ {% endif %}
+
  {{ run_hooks(post_hooks, inside_transaction=True) }}
 
  -- `COMMIT` happens here

core/dbt/include/global_project/macros/materializations/snapshot/snapshot.sql

@@ -263,6 +263,10 @@
 
  {% do persist_docs(target_relation, model) %}
 
+ {% if not target_relation_exists %}
+ {% do create_indexes(target_relation) %}
+ {% endif %}
+
  {{ run_hooks(post_hooks, inside_transaction=True) }}
 
  {{ adapter.commit() }}

core/dbt/include/global_project/macros/materializations/table/table.sql

@@ -47,6 +47,8 @@
 
  {{ adapter.rename_relation(intermediate_relation, target_relation) }}
 
+ {% do create_indexes(target_relation) %}
+
  {{ run_hooks(post_hooks, inside_transaction=True) }}
 
  {% do persist_docs(target_relation, model) %}

plugins/postgres/dbt/adapters/postgres/impl.py

@@ -1,21 +1,62 @@
+from datetime import datetime
 from dataclasses import dataclass
-from typing import Optional, Set
+from typing import Optional, Set, List, Any
 from dbt.adapters.base.meta import available
 from dbt.adapters.base.impl import AdapterConfig
 from dbt.adapters.sql import SQLAdapter
 from dbt.adapters.postgres import PostgresConnectionManager
 from dbt.adapters.postgres import PostgresColumn
 from dbt.adapters.postgres import PostgresRelation
+from dbt.dataclass_schema import dbtClassMixin, ValidationError
 import dbt.exceptions
+import dbt.utils
 
 
 # note that this isn't an adapter macro, so just a single underscore
 GET_RELATIONS_MACRO_NAME = 'postgres_get_relations'
 
 
+@dataclass
+class PostgresIndexConfig(dbtClassMixin):
+ columns: List[str]
+ unique: bool = False
+ type: Optional[str] = None
+
+ def render(self, relation):
+ # We append the current timestamp to the index name because otherwise
+ # the index will only be created on every other run. See
+ # https:/fishtown-analytics/dbt/issues/1945#issuecomment-576714925
+ # for an explanation.
+ now = datetime.utcnow().isoformat()
+ inputs = (self.columns +
+ [relation.render(), str(self.unique), str(self.type), now])
+ string = '_'.join(inputs)
+ return dbt.utils.md5(string)
+
+ @classmethod
+ def parse(cls, raw_index) -> Optional['PostgresIndexConfig']:
+ if raw_index is None:
+ return None
+ try:
+ cls.validate(raw_index)
+ return cls.from_dict(raw_index)
+ except ValidationError as exc:
+ msg = dbt.exceptions.validator_error_message(exc)
+ dbt.exceptions.raise_compiler_error(
+ f'Could not parse index config: {msg}'
+ )
+ except TypeError:
+ dbt.exceptions.raise_compiler_error(
+ f'Invalid index config:\n'
+ f' Got: {raw_index}\n'
+ f' Expected a dictionary with at minimum a "columns" key'
+ )
+
+
 @dataclass
 class PostgresConfig(AdapterConfig):
  unlogged: Optional[bool] = None
+ indexes: Optional[List[PostgresIndexConfig]] = None
 
 
 class PostgresAdapter(SQLAdapter):
@@ -42,6 +83,10 @@ def verify_database(self, database):
  # return an empty string on success so macros can call this
  return ''
 
+ @available
+ def parse_index(self, raw_index: Any) -> Optional[PostgresIndexConfig]:
+ return PostgresIndexConfig.parse(raw_index)
+
  def _link_cached_database_relations(self, schemas: Set[str]):
  """
  :param schemas: The set of schemas that should have links added.

plugins/postgres/dbt/include/postgres/macros/adapters.sql

@@ -14,6 +14,21 @@
  );
 {%- endmacro %}
 
+{% macro postgres__get_create_index_sql(relation, index_dict) -%}
+ {%- set index_config = adapter.parse_index(index_dict) -%}
+ {%- set comma_separated_columns = ", ".join(index_config.columns) -%}
+ {%- set index_name = index_config.render(relation) -%}
+
+ create {% if index_config.unique -%}
+ unique
+ {%- endif %} index if not exists
+ "{{ index_name }}"
+ on {{ relation }} {% if index_config.type -%}
+ using {{ index_config.type }}
+ {%- endif %}
+ ({{ comma_separated_columns }});
+{%- endmacro %}
+
 {% macro postgres__create_schema(relation) -%}
  {% if relation.database -%}
  {{ adapter.verify_database(relation.database) }}

test/integration/065_postgres_index_tests/data/seed.csv

@@ -0,0 +1,4 @@
+country_code,country_name
+US,United States
+CA,Canada
+GB,United Kingdom

test/integration/065_postgres_index_tests/models-invalid/invalid_columns_type.sql

@@ -0,0 +1,10 @@
+{{
+ config(
+ materialized = "table",
+ indexes=[
+ {'columns': 'column_a, column_b'},
+ ]
+ )
+}}
+
+select 1 as column_a, 2 as column_b

test/integration/065_postgres_index_tests/models-invalid/invalid_type.sql

@@ -0,0 +1,10 @@
+{{
+ config(
+ materialized = "table",
+ indexes=[
+ {'columns': ['column_a'], 'type': 'non_existent_type'},
+ ]
+ )
+}}
+
+select 1 as column_a, 2 as column_b

test/integration/065_postgres_index_tests/models-invalid/invalid_unique_config.sql

@@ -0,0 +1,10 @@
+{{
+ config(
+ materialized = "table",
+ indexes=[
+ {'columns': ['column_a'], 'unique': 'yes'},
+ ]
+ )
+}}
+
+select 1 as column_a, 2 as column_b

test/integration/065_postgres_index_tests/models-invalid/missing_columns.sql

@@ -0,0 +1,10 @@
+{{
+ config(
+ materialized = "table",
+ indexes=[
+ {'unique': True},
+ ]
+ )
+}}
+
+select 1 as column_a, 2 as column_b

test/integration/065_postgres_index_tests/models/incremental.sql

@@ -0,0 +1,18 @@
+{{
+ config(
+ materialized = "incremental",
+ indexes=[
+ {'columns': ['column_a'], 'type': 'hash'},
+ {'columns': ['column_a', 'column_b'], 'unique': True},
+ ]
+ )
+}}
+
+select *
+from (
+ select 1 as column_a, 2 as column_b
+) t
+
+{% if is_incremental() %}
+ where column_a > (select max(column_a) from {{this}})
+{% endif %}

test/integration/065_postgres_index_tests/models/table.sql

@@ -0,0 +1,14 @@
+{{
+ config(
+ materialized = "table",
+ indexes=[
+ {'columns': ['column_a']},
+ {'columns': ['column_b']},
+ {'columns': ['column_a', 'column_b']},
+ {'columns': ['column_b', 'column_a'], 'type': 'btree', 'unique': True},
+ {'columns': ['column_a'], 'type': 'hash'}
+ ]
+ )
+}}
+
+select 1 as column_a, 2 as column_b

test/integration/065_postgres_index_tests/snapshots/colors.sql

@@ -0,0 +1,29 @@
+{% snapshot colors %}
+
+ {{
+ config(
+ target_database=database,
+ target_schema=schema,
+ unique_key='id',
+ strategy='check',
+ check_cols=['color'],
+ indexes=[
+ {'columns': ['id'], 'type': 'hash'},
+ {'columns': ['id', 'color'], 'unique': True},
+ ]
+ )
+ }}
+
+ {% if var('version') == 1 %}
+
+ select 1 as id, 'red' as color union all
+ select 2 as id, 'green' as color
+
+ {% else %}
+
+ select 1 as id, 'blue' as color union all
+ select 2 as id, 'green' as color
+
+ {% endif %}
+
+{% endsnapshot %}