Merge pull request #22 from l-gysin/main

Not Null Constraints for not_null and primary_key tests

dbt_project.yml

@@ -33,3 +33,4 @@ vars:
   dbt_constraints_sources_pk_enabled: true
   dbt_constraints_sources_uk_enabled: true
   dbt_constraints_sources_fk_enabled: true
+  dbt_constraints_sources_nn_enabled: true

integration_tests/dbt_project.yml

@@ -48,3 +48,4 @@ models:
 
 seeds:
   +quote_columns: false
+  +full_refresh: false

macros/create_constraints.sql

@@ -65,6 +65,10 @@
 {%- endmacro -%}
 
 
+{%- macro create_not_null(table_model, column_names, verify_permissions, quote_columns=false) -%}
+    {{ return(adapter.dispatch('create_not_null', 'dbt_constraints')(table_model, column_names, verify_permissions, quote_columns)) }}
+{%- endmacro -%}
+
 
 {#- Define two macros for detecting if PK, UK, and FK exist that can be overridden by DB implementations -#}
 
@@ -95,6 +99,11 @@
   {{ return(adapter.dispatch('truncate_relation')(relation)) }}
 {% endmacro %}
 
+{#- Override dbt's drop_relation macro to allow us to create adapter specific versions that drop constraints -#}
+
+{% macro drop_relation(relation) -%}
+  {{ return(adapter.dispatch('drop_relation')(relation)) }}
+{% endmacro %}
 
 
 
@@ -110,7 +119,8 @@
             'unique_combination_of_columns',
             'unique',
             'foreign_key',
-            'relationships'],
+            'relationships',
+            'not_null'],
         quote_columns=false) -%}
     {%- if execute and var('dbt_constraints_enabled', false) -%}
         {%- do log("Running dbt Constraints", info=true) -%}
@@ -133,6 +143,9 @@
         {%- if 'relationships' in constraint_types -%}
             {%- do dbt_constraints.create_constraints_by_type(['relationships'], quote_columns) -%}
         {%- endif -%}
+        {%- if 'not_null' in constraint_types -%}
+            {%- do dbt_constraints.create_constraints_by_type(['not_null'], quote_columns) -%}
+        {%- endif -%}
 
         {%- do log("Finished dbt Constraints", info=true) -%}
     {%- endif -%}
@@ -162,7 +175,6 @@
             These models must be physical tables and cannot be sources -#}
         {%- set table_models = [] -%}
         {%- for node in graph.nodes.values() | selectattr("unique_id", "in", test_model.depends_on.nodes)
-                {#- Update to include snapshot resource type -#}
                 if node.resource_type in ( ( "model", "snapshot") )
                     if node.config.materialized in( ("table", "incremental", "snapshot") ) -%}
 
@@ -176,6 +188,7 @@
                 ( var('dbt_constraints_sources_pk_enabled', false) and test_model.test_metadata.name in("primary_key") )
              or ( var('dbt_constraints_sources_uk_enabled', false) and test_model.test_metadata.name in("unique_key", "unique_combination_of_columns", "unique") )
              or ( var('dbt_constraints_sources_fk_enabled', false) and test_model.test_metadata.name in("foreign_key", "relationships") )
+             or ( var('dbt_constraints_sources_nn_enabled', false) and test_model.test_metadata.name in("not_null") )
             ) -%}
             {%- for node in graph.sources.values()
                 | selectattr("resource_type", "equalto", "source")
@@ -216,6 +229,7 @@
                 identifier=table_models[0].alias ) -%}
             {%- if dbt_constraints.table_columns_all_exist(table_relation, column_names) -%}
                 {%- if test_model.test_metadata.name == "primary_key" -%}
+                    {%- do dbt_constraints.create_not_null(table_relation, column_names, ns.verify_permissions, quote_columns) -%}
                     {%- do dbt_constraints.create_primary_key(table_relation, column_names, ns.verify_permissions, quote_columns) -%}
                 {%- else  -%}
                     {%- do dbt_constraints.create_unique_key(table_relation, column_names, ns.verify_permissions, quote_columns) -%}
@@ -297,6 +311,37 @@
             {%- else  -%}
                 {%- do log("Skipping foreign key because a we couldn't find the child table: model=" ~ fk_model_names ~ " or source=" ~ fk_source_names, info=true) -%}
             {%- endif -%}
+
+        {#- We only create NN if there is one model referenced by the test
+            and if all the columns exist as physical columns on the table -#}
+        {%- elif 1 == table_models|count
+            and test_model.test_metadata.name in("not_null") -%}
+
+            {# Attempt to identify a parameter we can use for the column names #}
+            {%- set column_names = [] -%}
+            {%- if  test_parameters.column_names -%}
+                {%- set column_names =  test_parameters.column_names -%}
+            {%- elif  test_parameters.combination_of_columns -%}
+                {%- set column_names =  test_parameters.combination_of_columns -%}
+            {%- elif  test_parameters.column_name -%}
+                {%- set column_names =  [test_parameters.column_name] -%}
+            {%- else  -%}
+                {{ exceptions.raise_compiler_error(
+                "`column_names` or `column_name` parameter missing for not null constraint on table: '" ~ table_models[0].name
+                ) }}
+            {%- endif -%}
+
+            {%- set table_relation = api.Relation.create(
+                database=table_models[0].database,
+                schema=table_models[0].schema,
+                identifier=table_models[0].alias ) -%}
+
+            {%- if dbt_constraints.table_columns_all_exist(table_relation, column_names) -%}
+                {%- do dbt_constraints.create_not_null(table_relation, column_names, ns.verify_permissions, quote_columns) -%}
+            {%- else  -%}
+                {%- do log("Skipping not null constraint because a physical column name was not found on the table: " ~ table_models[0].name ~ " " ~ column_names, info=true) -%}
+            {%- endif -%}
+
         {%- endif -%}
 
     {%- endfor -%}
@@ -339,7 +384,3 @@
         {{ return(false) }}
     {%- endif -%}
 {%- endmacro -%}
-
-
-{%- macro drop_constraints() -%}
-{%- endmacro -%}

macros/oracle__create_constraints.sql

@@ -130,8 +130,33 @@ END;
 
 {%- endmacro -%}
 
+{# Oracle specific implementation to create a not null constraint #}
+{%- macro oracle__create_not_null(table_relation, column_names, verify_permissions, quote_columns=false) -%}
+    {%- set columns_list = dbt_constraints.get_quoted_column_list(column_names, quote_columns) -%}
 
+    {%- if dbt_constraints.have_ownership_priv(table_relation, verify_permissions) -%}
 
+            {%- set modify_statements= [] -%}
+            {%- for column in columns_list -%}
+                {%- set modify_statements = modify_statements.append( column ~ " NOT NULL" ) -%}
+            {%- endfor -%}
+            {%- set modify_statement_csv = modify_statements | join(", ") -%}
+            {%- set query -%}
+BEGIN
+    EXECUTE IMMEDIATE 'ALTER TABLE {{table_relation}} MODIFY ( {{ modify_statement_csv }} )';
+EXCEPTION
+    WHEN OTHERS THEN
+        DBMS_OUTPUT.ENABLE(BUFFER_SIZE => NULL);
+        DBMS_OUTPUT.PUT_LINE('Unable to create constraint: ' || SQLERRM);
+END;
+            {%- endset -%}
+            {%- do log("Creating not null constraint for: " ~ columns_list | join(", ") ~ " in " ~ table_relation, info=true) -%}
+            {%- do run_query(query) -%}
+
+    {%- else -%}
+        {%- do log("Skipping not null constraint for " ~ columns_list | join(", ") ~ " in " ~ table_relation ~ " because of insufficient privileges: " ~ table_relation, info=true) -%}
+    {%- endif -%}
+{%- endmacro -%}
 {#- This macro is used in create macros to avoid duplicate PK/UK constraints
     and to skip FK where no PK/UK constraint exists on the parent table -#}
 {%- macro oracle__unique_constraint_exists(table_relation, column_names) -%}
@@ -208,3 +233,41 @@ order by 1, 2
 {%- macro oracle__have_ownership_priv(table_relation, verify_permissions) -%}
     {{ return(true) }}
 {%- endmacro -%}
+
+{% macro oracle__drop_referential_constraints(relation) -%}
+    {%- call statement('drop_constraint_cascade') -%}
+BEGIN
+    FOR REC IN (
+        SELECT constraint_name
+        FROM all_constraints cons
+        WHERE cons.constraint_type IN ('P', 'U', 'R')
+            AND upper(cons.owner) = upper('{{relation.schema}}')
+            AND upper(cons.table_name) = upper('{{relation.identifier}}')
+        ORDER BY 1
+    ) LOOP
+        BEGIN
+            EXECUTE IMMEDIATE 'ALTER TABLE {{relation}} DROP CONSTRAINT "'||REC.CONSTRAINT_NAME||'" CASCADE';
+        EXCEPTION
+            WHEN OTHERS THEN
+                DBMS_OUTPUT.ENABLE(BUFFER_SIZE => NULL);
+                DBMS_OUTPUT.PUT_LINE('Unable to drop constraint: ' || SQLERRM);
+        END;
+    END LOOP;
+END;
+    {%- endcall -%}
+
+{% endmacro %}
+
+{#- Oracle will error if you try to truncate tables with FK constraints or tables with PK/UK constraints
+    referenced by FK so we will drop all constraints before truncating tables -#}
+{% macro oracle__truncate_relation(relation) -%}
+    {{ oracle__drop_referential_constraints(relation) }}
+    {{ return(adapter.dispatch('truncate_relation', 'dbt')(relation)) }}
+{% endmacro %}
+
+{#- Oracle will error if you try to drop tables with FK constraints or tables with PK/UK constraints
+    referenced by FK so we will drop all constraints before dropping tables -#}
+{% macro oracle__drop_relation(relation) -%}
+    {{ oracle__drop_referential_constraints(relation) }}
+    {{ return(adapter.dispatch('drop_relation', 'dbt')(relation)) }}
+{% endmacro %}

macros/postgres__create_constraints.sql

@@ -17,11 +17,10 @@
 
         {%- if dbt_constraints.have_ownership_priv(table_relation, verify_permissions) -%}
 
-            {%- set query -%}
-            ALTER TABLE {{table_relation}} ADD CONSTRAINT {{constraint_name}} PRIMARY KEY ( {{columns_csv}} )
-            {%- endset -%}
             {%- do log("Creating primary key: " ~ constraint_name, info=true) -%}
-            {%- do run_query(query) -%}
+            {%- call statement('add_pk', fetch_result=False, auto_begin=True) -%}
+            ALTER TABLE {{table_relation}} ADD CONSTRAINT {{constraint_name}} PRIMARY KEY ( {{columns_csv}} )
+            {%- endcall -%}
             {{ adapter.commit() }}
 
         {%- else -%}
@@ -55,11 +54,10 @@
 
         {%- if dbt_constraints.have_ownership_priv(table_relation, verify_permissions) -%}
 
-            {%- set query -%}
-            ALTER TABLE {{table_relation}} ADD CONSTRAINT {{constraint_name}} UNIQUE ( {{columns_csv}} )
-            {%- endset -%}
             {%- do log("Creating unique key: " ~ constraint_name, info=true) -%}
-            {%- do run_query(query) -%}
+            {%- call statement('add_uk', fetch_result=False, auto_begin=True) -%}
+            ALTER TABLE {{table_relation}} ADD CONSTRAINT {{constraint_name}} UNIQUE ( {{columns_csv}} )
+            {%- endcall -%}
             {{ adapter.commit() }}
 
         {%- else -%}
@@ -72,7 +70,27 @@
 
 {%- endmacro -%}
 
+{# PostgreSQL specific implementation to create a not null constraint #}
+{%- macro postgres__create_not_null(table_relation, column_names, verify_permissions, quote_columns=false) -%}
+    {%- set columns_list = dbt_constraints.get_quoted_column_list(column_names, quote_columns) -%}
+
+    {%- if dbt_constraints.have_ownership_priv(table_relation, verify_permissions) -%}
+
+            {%- set modify_statements= [] -%}
+            {%- for column in columns_list -%}
+                {%- set modify_statements = modify_statements.append( "ALTER COLUMN " ~ column ~ " SET NOT NULL" ) -%}
+            {%- endfor -%}
+            {%- set modify_statement_csv = modify_statements | join(", ") -%}
+            {%- do log("Creating not null constraint for: " ~ columns_list | join(", ") ~ " in " ~ table_relation, info=true) -%}
+            {%- call statement('add_nn', fetch_result=False, auto_begin=True) -%}
+                ALTER TABLE {{table_relation}} {{ modify_statement_csv }};
+            {%- endcall -%}
+            {{ adapter.commit() }}
 
+    {%- else -%}
+        {%- do log("Skipping not null constraint for " ~ columns_list | join(", ") ~ " in " ~ table_relation ~ " because of insufficient privileges: " ~ table_relation, info=true) -%}
+    {%- endif -%}
+{%- endmacro -%}
 
 {# PostgreSQL specific implementation to create a foreign key #}
 {%- macro postgres__create_foreign_key(pk_table_relation, pk_column_names, fk_table_relation, fk_column_names, verify_permissions, quote_columns=true) -%}
@@ -95,11 +113,10 @@
 
             {%- if dbt_constraints.have_ownership_priv(fk_table_relation, verify_permissions) and dbt_constraints.have_references_priv(pk_table_relation, verify_permissions) -%}
 
-                {%- set query -%}
-                ALTER TABLE {{fk_table_relation}} ADD CONSTRAINT {{constraint_name}} FOREIGN KEY ( {{fk_columns_csv}} ) REFERENCES {{pk_table_relation}} ( {{pk_columns_csv}} ) ON DELETE NO ACTION DEFERRABLE INITIALLY DEFERRED
-                {%- endset -%}
                 {%- do log("Creating foreign key: " ~ constraint_name ~ " referencing " ~ pk_table_relation.identifier ~ " " ~ pk_column_names, info=true) -%}
-                {%- do run_query(query) -%}
+                {%- call statement('add_fk', fetch_result=False, auto_begin=True) -%}
+                ALTER TABLE {{fk_table_relation}} ADD CONSTRAINT {{constraint_name}} FOREIGN KEY ( {{fk_columns_csv}} ) REFERENCES {{pk_table_relation}} ( {{pk_columns_csv}} ) ON DELETE NO ACTION DEFERRABLE INITIALLY DEFERRED
+                {%- endcall -%}
                 {{ adapter.commit() }}
 
             {%- else -%}
@@ -231,11 +248,7 @@
 {%- endmacro -%}
 
 
-
-
-{#- PostgreSQL will error if you try to truncate tables with FK constraints or tables with PK/UK constraints
-    referenced by FK so we will drop all constraints before truncating tables -#}
-{% macro postgres__truncate_relation(relation) -%}
+{% macro postgres__drop_referential_constraints(relation) -%}
     {%- set lookup_query -%}
     select constraint_name
     from information_schema.table_constraints
@@ -246,12 +259,25 @@
     {%- set constraint_list = run_query(lookup_query) -%}
 
     {%- for constraint_name in constraint_list.columns["constraint_name"].values() -%}
-        {%- set drop_statement -%}
-        ALTER TABLE {{relation}} DROP CONSTRAINT "{{constraint_name}}" CASCADE
-        {%- endset -%}
         {%- do log("Dropping constraint: " ~ constraint_name ~ " from table " ~ relation, info=false) -%}
-        {%- do run_query(drop_statement) -%}
+        {%- call statement('drop_constraint_cascade', fetch_result=False, auto_begin=True) -%}
+        ALTER TABLE {{relation}} DROP CONSTRAINT IF EXISTS "{{constraint_name}}" CASCADE
+        {%- endcall -%}
+        {{ adapter.commit() }}
     {% endfor %}
 
+{% endmacro %}
+
+{#- PostgreSQL will error if you try to truncate tables with FK constraints or tables with PK/UK constraints
+    referenced by FK so we will drop all constraints before truncating tables -#}
+{% macro postgres__truncate_relation(relation) -%}
+    {{ postgres__drop_referential_constraints(relation) }}
     {{ return(adapter.dispatch('truncate_relation', 'dbt')(relation)) }}
 {% endmacro %}
+
+{#- PostgreSQL will get deadlocks if you try to drop tables with FK constraints or tables with PK/UK constraints
+    referenced by FK so we will drop all constraints before dropping tables -#}
+{% macro postgres__drop_relation(relation) -%}
+    {{ postgres__drop_referential_constraints(relation) }}
+    {{ return(adapter.dispatch('drop_relation', 'dbt')(relation)) }}
+{% endmacro %}

macros/snowflake__create_constraints.sql

@@ -86,7 +86,27 @@
 
 {%- endmacro -%}
 
+{# Snowflake specific implementation to create a not null constraint #}
+{%- macro snowflake__create_not_null(table_relation, column_names, verify_permissions, quote_columns=false) -%}
+    {%- set columns_list = dbt_constraints.get_quoted_column_list(column_names, quote_columns) -%}
 
+    {%- if dbt_constraints.have_ownership_priv(table_relation, verify_permissions) -%}
+
+            {%- set modify_statements= [] -%}
+            {%- for column in columns_list -%}
+                {%- set modify_statements = modify_statements.append( "COLUMN " ~ column ~ " SET NOT NULL" ) -%}
+            {%- endfor -%}
+            {%- set modify_statement_csv = modify_statements | join(", ") -%}
+            {%- set query -%}
+                ALTER TABLE {{table_relation}} MODIFY {{ modify_statement_csv }};
+            {%- endset -%}
+            {%- do log("Creating not null constraint for: " ~ columns_list | join(", ") ~ " in " ~ table_relation, info=true) -%}
+            {%- do run_query(query) -%}
+
+    {%- else -%}
+        {%- do log("Skipping not null constraint for " ~ columns_list | join(", ") ~ " in " ~ table_relation ~ " because of insufficient privileges: " ~ table_relation, info=true) -%}
+    {%- endif -%}
+{%- endmacro -%}
 
 {#- This macro is used in create macros to avoid duplicate PK/UK constraints
     and to skip FK where no PK/UK constraint exists on the parent table -#}