Merge pull request #382 from dbt-msft/patch-1

Add schema authorization config

Author: sdebruyn

dbt/adapters/sqlserver/sql_server_adapter.py

@@ -2,7 +2,11 @@
 
 import agate
 from dbt.adapters.base.relation import BaseRelation
+from dbt.adapters.cache import _make_ref_key_msg
 from dbt.adapters.sql import SQLAdapter
+from dbt.adapters.sql.impl import CREATE_SCHEMA_MACRO_NAME
+from dbt.events.functions import fire_event
+from dbt.events.types import SchemaCreation
 
 from dbt.adapters.sqlserver.sql_server_column import SQLServerColumn
 from dbt.adapters.sqlserver.sql_server_configs import SQLServerConfigs
@@ -14,6 +18,21 @@ class SQLServerAdapter(SQLAdapter):
     Column = SQLServerColumn
     AdapterSpecificConfigs = SQLServerConfigs
 
+    def create_schema(self, relation: BaseRelation) -> None:
+        relation = relation.without_identifier()
+        fire_event(SchemaCreation(relation=_make_ref_key_msg(relation)))
+        macro_name = CREATE_SCHEMA_MACRO_NAME
+        kwargs = {
+            "relation": relation,
+        }
+
+        if self.config.credentials.schema_authorization:
+            kwargs["schema_authorization"] = self.config.credentials.schema_authorization
+            macro_name = "sqlserver__create_schema_with_authorization"
+
+        self.execute_macro(macro_name, kwargs=kwargs)
+        self.commit_if_has_connection()
+
     @classmethod
     def date_function(cls):
         return "getdate()"

dbt/adapters/sqlserver/sql_server_credentials.py

@@ -21,6 +21,7 @@ class SQLServerCredentials(Credentials):
     encrypt: Optional[bool] = True  # default value in MS ODBC Driver 18 as well
     trust_cert: Optional[bool] = False  # default value in MS ODBC Driver 18 as well
     retries: int = 1
+    schema_authorization: Optional[str] = None
 
     _ALIASES = {
         "user": "UID",
@@ -33,6 +34,7 @@ class SQLServerCredentials(Credentials):
         "app_id": "client_id",
         "app_secret": "client_secret",
         "TrustServerCertificate": "trust_cert",
+        "schema_auth": "schema_authorization",
     }
 
     @property

dbt/include/sqlserver/macros/adapters/schema.sql

@@ -1,9 +1,19 @@
 {% macro sqlserver__create_schema(relation) -%}
   {% call statement('create_schema') -%}
     USE [{{ relation.database }}];
-    IF NOT EXISTS (SELECT * FROM sys.schemas WHERE name = '{{ relation.without_identifier().schema }}')
+    IF NOT EXISTS (SELECT * FROM sys.schemas WHERE name = '{{ relation.schema }}')
     BEGIN
-    EXEC('CREATE SCHEMA [{{ relation.without_identifier().schema }}]')
+    EXEC('CREATE SCHEMA [{{ relation.schema }}]')
+    END
+  {% endcall %}
+{% endmacro %}
+
+{% macro sqlserver__create_schema_with_authorization(relation, schema_authorization) -%}
+  {% call statement('create_schema') -%}
+    USE [{{ relation.database }}];
+    IF NOT EXISTS (SELECT * FROM sys.schemas WHERE name = '{{ relation.schema }}')
+    BEGIN
+    EXEC('CREATE SCHEMA [{{ relation.schema }}] AUTHORIZATION [{{ schema_authorization }}]')
     END
   {% endcall %}
 {% endmacro %}

tests/conftest.py

@@ -13,25 +13,33 @@ def pytest_addoption(parser):
 
 
 @pytest.fixture(scope="class")
-def dbt_profile_target(request: FixtureRequest):
+def dbt_profile_target(request: FixtureRequest, dbt_profile_target_update):
     profile = request.config.getoption("--profile")
 
     if profile == "ci_sql_server":
-        return _profile_ci_sql_server()
-    if profile == "ci_azure_cli":
-        return _profile_ci_azure_cli()
-    if profile == "ci_azure_auto":
-        return _profile_ci_azure_auto()
-    if profile == "ci_azure_environment":
-        return _profile_ci_azure_environment()
-    if profile == "ci_azure_basic":
-        return _profile_ci_azure_basic()
-    if profile == "user":
-        return _profile_user()
-    if profile == "user_azure":
-        return _profile_user_azure()
-
-    raise ValueError(f"Unknown profile: {profile}")
+        target = _profile_ci_sql_server()
+    elif profile == "ci_azure_cli":
+        target = _profile_ci_azure_cli()
+    elif profile == "ci_azure_auto":
+        target = _profile_ci_azure_auto()
+    elif profile == "ci_azure_environment":
+        target = _profile_ci_azure_environment()
+    elif profile == "ci_azure_basic":
+        target = _profile_ci_azure_basic()
+    elif profile == "user":
+        target = _profile_user()
+    elif profile == "user_azure":
+        target = _profile_user_azure()
+    else:
+        raise ValueError(f"Unknown profile: {profile}")
+
+    target.update(dbt_profile_target_update)
+    return target
+
+
+@pytest.fixture(scope="class")
+def dbt_profile_target_update():
+    return {}
 
 
 @pytest.fixture(scope="class")

tests/functional/adapter/test_schema.py

@@ -0,0 +1,37 @@
+import os
+
+import pytest
+from dbt.tests.util import run_dbt
+
+
+class TestSchemaCreation:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "dummy.sql": """
+{{ config(schema='with_custom_auth') }}
+select 1 as id
+""",
+        }
+
+    @staticmethod
+    @pytest.fixture(scope="class")
+    def dbt_profile_target_update():
+        return {"schema_authorization": "{{ env_var('DBT_TEST_USER_1') }}"}
+
+    @staticmethod
+    def _verify_schema_owner(schema_name, owner, project):
+        get_schema_owner = f"""
+select schema_owner from information_schema.schemata where schema_name = '{schema_name}'
+        """
+        result = project.run_sql(get_schema_owner, fetch="one")[0]
+        assert result == owner
+
+    def test_schema_creation(self, project, unique_schema):
+        res = run_dbt(["run"])
+        assert len(res) == 1
+
+        self._verify_schema_owner(unique_schema, os.getenv("DBT_TEST_USER_1"), project)
+        self._verify_schema_owner(
+            unique_schema + "_with_custom_auth", os.getenv("DBT_TEST_USER_1"), project
+        )