-
Notifications
You must be signed in to change notification settings - Fork 76
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add support for model contracts in v1.5 #148
Merged
jwills
merged 14 commits into
duckdb:jwills_rc150
from
jtcohen6:jerco/v1.5-model-contracts
Apr 26, 2023
Merged
Changes from all commits
Commits
Show all changes
14 commits
Select commit
Hold shift + click to select a range
b2c74d1
Add support for model contracts in v1.5
jtcohen6 6626430
Add references alias for foreign_key
jtcohen6 64f351d
Fixup create_table_as logic
jtcohen6 b5b488e
Add back the test imports I foolishly removed
jtcohen6 32a9dfa
Reorder logic, fix tests
jtcohen6 40e56f1
First up: this is a bit nicer
jwills 6f70104
Just trying to see where this leads me
jwills 2aeb1ff
fix that
jwills 6fd767c
stopping here for the night; need to do most of the localenv stuff tmrw
jwills 2465dba
Renaming/refactoring some things
jwills 7d81bd9
now with format fixes
jwills 8d5b133
let's try whistling this
jwills 259fd3b
simplify this back down
jwills 5a76e45
add in tz-aware timestamp test
jwills File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,95 @@ | ||
from . import Environment | ||
from .. import credentials | ||
from .. import utils | ||
from dbt.contracts.connection import AdapterResponse | ||
from dbt.exceptions import DbtRuntimeError | ||
|
||
|
||
class DuckDBCursorWrapper: | ||
def __init__(self, cursor): | ||
self._cursor = cursor | ||
|
||
# forward along all non-execute() methods/attribute look-ups | ||
def __getattr__(self, name): | ||
return getattr(self._cursor, name) | ||
|
||
def execute(self, sql, bindings=None): | ||
try: | ||
if bindings is None: | ||
return self._cursor.execute(sql) | ||
else: | ||
return self._cursor.execute(sql, bindings) | ||
except RuntimeError as e: | ||
raise DbtRuntimeError(str(e)) | ||
|
||
|
||
class DuckDBConnectionWrapper: | ||
def __init__(self, cursor): | ||
self._cursor = DuckDBCursorWrapper(cursor) | ||
|
||
def close(self): | ||
self._cursor.close() | ||
|
||
def cursor(self): | ||
return self._cursor | ||
|
||
|
||
class LocalEnvironment(Environment): | ||
def __init__(self, credentials: credentials.DuckDBCredentials): | ||
self.conn = self.initialize_db(credentials) | ||
self._plugins = self.initialize_plugins(credentials) | ||
self.creds = credentials | ||
|
||
def handle(self): | ||
# Extensions/settings need to be configured per cursor | ||
cursor = self.initialize_cursor(self.creds, self.conn.cursor()) | ||
return DuckDBConnectionWrapper(cursor) | ||
|
||
def submit_python_job(self, handle, parsed_model: dict, compiled_code: str) -> AdapterResponse: | ||
con = handle.cursor() | ||
|
||
def ldf(table_name): | ||
return con.query(f"select * from {table_name}") | ||
|
||
self.run_python_job(con, ldf, parsed_model["alias"], compiled_code) | ||
return AdapterResponse(_message="OK") | ||
|
||
def load_source(self, plugin_name: str, source_config: utils.SourceConfig): | ||
if plugin_name not in self._plugins: | ||
raise Exception( | ||
f"Plugin {plugin_name} not found; known plugins are: " | ||
+ ",".join(self._plugins.keys()) | ||
) | ||
plugin = self._plugins[plugin_name] | ||
handle = self.handle() | ||
cursor = handle.cursor() | ||
save_mode = source_config.meta.get("save_mode", "overwrite") | ||
if save_mode in ("ignore", "error_if_exists"): | ||
schema, identifier = source_config.schema, source_config.identifier | ||
q = f"""SELECT COUNT(1) | ||
FROM information_schema.tables | ||
WHERE table_schema = '{schema}' | ||
AND table_name = '{identifier}' | ||
""" | ||
if cursor.execute(q).fetchone()[0]: | ||
if save_mode == "error_if_exists": | ||
raise Exception(f"Source {source_config.table_name()} already exists!") | ||
else: | ||
# Nothing to do (we ignore the existing table) | ||
return | ||
df = plugin.load(source_config) | ||
assert df is not None | ||
materialization = source_config.meta.get("materialization", "table") | ||
cursor.execute( | ||
f"CREATE OR REPLACE {materialization} {source_config.table_name()} AS SELECT * FROM df" | ||
) | ||
cursor.close() | ||
handle.close() | ||
|
||
def close(self): | ||
if self.conn: | ||
self.conn.close() | ||
self.conn = None | ||
|
||
def __del__(self): | ||
self.close() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -8,12 +8,15 @@ | |
|
||
from dbt.adapters.base import BaseRelation | ||
from dbt.adapters.base.column import Column | ||
from dbt.adapters.base.impl import ConstraintSupport | ||
from dbt.adapters.base.meta import available | ||
from dbt.adapters.duckdb.connections import DuckDBConnectionManager | ||
from dbt.adapters.duckdb.glue import create_or_update_table | ||
from dbt.adapters.duckdb.relation import DuckDBRelation | ||
from dbt.adapters.sql import SQLAdapter | ||
from dbt.contracts.connection import AdapterResponse | ||
from dbt.contracts.graph.nodes import ColumnLevelConstraint | ||
from dbt.contracts.graph.nodes import ConstraintType | ||
from dbt.exceptions import DbtInternalError | ||
from dbt.exceptions import DbtRuntimeError | ||
|
||
|
@@ -22,6 +25,14 @@ class DuckDBAdapter(SQLAdapter): | |
ConnectionManager = DuckDBConnectionManager | ||
Relation = DuckDBRelation | ||
|
||
CONSTRAINT_SUPPORT = { | ||
ConstraintType.check: ConstraintSupport.ENFORCED, | ||
ConstraintType.not_null: ConstraintSupport.ENFORCED, | ||
ConstraintType.unique: ConstraintSupport.ENFORCED, | ||
ConstraintType.primary_key: ConstraintSupport.ENFORCED, | ||
ConstraintType.foreign_key: ConstraintSupport.ENFORCED, | ||
} | ||
|
||
@classmethod | ||
def date_function(cls) -> str: | ||
return "now()" | ||
|
@@ -176,6 +187,30 @@ def get_rows_different_sql( | |
) | ||
return sql | ||
|
||
@available.parse(lambda *a, **k: []) | ||
def get_column_schema_from_query(self, sql: str) -> List[Column]: | ||
"""Get a list of the Columns with names and data types from the given sql.""" | ||
|
||
# Taking advantage of yet another amazing DuckDB SQL feature right here: the | ||
# ability to DESCRIBE a query instead of a relation | ||
describe_sql = f"DESCRIBE ({sql})" | ||
_, cursor = self.connections.add_select_query(describe_sql) | ||
ret = [] | ||
for row in cursor.fetchall(): | ||
name, dtype = row[0], row[1] | ||
ret.append(Column.create(name, dtype)) | ||
return ret | ||
|
||
@classmethod | ||
def render_column_constraint(cls, constraint: ColumnLevelConstraint) -> Optional[str]: | ||
"""Render the given constraint as DDL text. Should be overriden by adapters which need custom constraint | ||
rendering.""" | ||
if constraint.type == ConstraintType.foreign_key: | ||
# DuckDB doesn't support 'foreign key' as an alias | ||
return f"references {constraint.expression}" | ||
Comment on lines
+209
to
+210
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Actually
The majority of columnar db's don't enforce them, so it hasn't felt like a priority. It's neat that DuckDB actually does. |
||
else: | ||
return super().render_column_constraint(constraint) | ||
|
||
|
||
# Change `table_a/b` to `table_aaaaa/bbbbb` to avoid duckdb binding issues when relation_a/b | ||
# is called "table_a" or "table_b" in some of the dbt tests | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@jtcohen6 after a truly circuitous journey, this ended up being simple and delightful to implement (and in an environment-independent way to boot!)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
whoa! this is very handy
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
this is pretty sweet