From c399574a0ce6026bd29374845fd5320d4ba75252 Mon Sep 17 00:00:00 2001 From: Nathan <148575555+nathan-artie@users.noreply.github.com> Date: Mon, 9 Dec 2024 13:29:57 -0800 Subject: [PATCH] Add truncate table and drop table to dialects (#1071) --- clients/bigquery/dialect/ddl.go | 8 ++++++ clients/bigquery/dialect/dialect_test.go | 14 ++++++++++ clients/databricks/dialect/ddl.go | 8 ++++++ clients/databricks/dialect/dialect_test.go | 14 ++++++++++ clients/mssql/dialect/ddl.go | 8 ++++++ clients/mssql/dialect/dialect_test.go | 14 ++++++++++ clients/redshift/dialect/ddl.go | 30 ++++++++++++++-------- clients/redshift/dialect/dialect_test.go | 14 ++++++++++ clients/snowflake/dialect/ddl.go | 8 ++++++ clients/snowflake/dialect/dialect_test.go | 14 ++++++++++ lib/sql/dialect.go | 2 ++ 11 files changed, 123 insertions(+), 11 deletions(-) diff --git a/clients/bigquery/dialect/ddl.go b/clients/bigquery/dialect/ddl.go index 41df8ef88..fd2cb932f 100644 --- a/clients/bigquery/dialect/ddl.go +++ b/clients/bigquery/dialect/ddl.go @@ -24,6 +24,14 @@ func (BigQueryDialect) BuildCreateTableQuery(tableID sql.TableIdentifier, tempor } } +func (BigQueryDialect) BuildDropTableQuery(tableID sql.TableIdentifier) string { + return "DROP TABLE IF EXISTS " + tableID.FullyQualifiedName() +} + +func (BigQueryDialect) BuildTruncateTableQuery(tableID sql.TableIdentifier) string { + return "TRUNCATE TABLE " + tableID.FullyQualifiedName() +} + func (BigQueryDialect) BuildAddColumnQuery(tableID sql.TableIdentifier, sqlPart string) string { return fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s", tableID.FullyQualifiedName(), sqlPart) } diff --git a/clients/bigquery/dialect/dialect_test.go b/clients/bigquery/dialect/dialect_test.go index 06dbce1ee..d9df44f9f 100644 --- a/clients/bigquery/dialect/dialect_test.go +++ b/clients/bigquery/dialect/dialect_test.go @@ -66,6 +66,20 @@ func TestBigQueryDialect_BuildCreateTableQuery(t *testing.T) { ) } +func TestBigQueryDialect_BuildDropTableQuery(t *testing.T) { + assert.Equal(t, + "DROP TABLE IF EXISTS `project1`.`dataset2`.`table3`", + BigQueryDialect{}.BuildDropTableQuery(NewTableIdentifier("project1", "dataset2", "table3")), + ) +} + +func TestBigQueryDialect_BuildTruncateTableQuery(t *testing.T) { + assert.Equal(t, + "TRUNCATE TABLE `project1`.`dataset2`.`table3`", + BigQueryDialect{}.BuildTruncateTableQuery(NewTableIdentifier("project1", "dataset2", "table3")), + ) +} + func TestBigQueryDialect_BuildDropColumnQuery(t *testing.T) { fakeTableID := &mocks.FakeTableIdentifier{} fakeTableID.FullyQualifiedNameReturns("{TABLE}") diff --git a/clients/databricks/dialect/ddl.go b/clients/databricks/dialect/ddl.go index f38429022..e910a201c 100644 --- a/clients/databricks/dialect/ddl.go +++ b/clients/databricks/dialect/ddl.go @@ -12,6 +12,14 @@ func (DatabricksDialect) BuildCreateTableQuery(tableID sql.TableIdentifier, _ bo return fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (%s)", tableID.FullyQualifiedName(), strings.Join(colSQLParts, ", ")) } +func (DatabricksDialect) BuildDropTableQuery(tableID sql.TableIdentifier) string { + return "DROP TABLE IF EXISTS " + tableID.FullyQualifiedName() +} + +func (DatabricksDialect) BuildTruncateTableQuery(tableID sql.TableIdentifier) string { + return "TRUNCATE TABLE " + tableID.FullyQualifiedName() +} + func (DatabricksDialect) BuildAddColumnQuery(tableID sql.TableIdentifier, sqlPart string) string { return fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s", tableID.FullyQualifiedName(), sqlPart) } diff --git a/clients/databricks/dialect/dialect_test.go b/clients/databricks/dialect/dialect_test.go index 59c0c6b4d..6d3c3fc59 100644 --- a/clients/databricks/dialect/dialect_test.go +++ b/clients/databricks/dialect/dialect_test.go @@ -70,6 +70,20 @@ func TestDatabricksDialect_BuildCreateTableQuery(t *testing.T) { } } +func TestDatabricksDialect_BuildDropTableQuery(t *testing.T) { + assert.Equal(t, + "DROP TABLE IF EXISTS `project1`.`dataset2`.`table3`", + DatabricksDialect{}.BuildDropTableQuery(NewTableIdentifier("project1", "dataset2", "table3")), + ) +} + +func TestDatabricksDialect_BuildTruncateTableQuery(t *testing.T) { + assert.Equal(t, + "TRUNCATE TABLE `project1`.`dataset2`.`table3`", + DatabricksDialect{}.BuildTruncateTableQuery(NewTableIdentifier("project1", "dataset2", "table3")), + ) +} + func TestDatabricksDialect_BuildAddColumnQuery(t *testing.T) { fakeTableID := &mocks.FakeTableIdentifier{} fakeTableID.FullyQualifiedNameReturns("{TABLE}") diff --git a/clients/mssql/dialect/ddl.go b/clients/mssql/dialect/ddl.go index c810239f3..cf84cc1d1 100644 --- a/clients/mssql/dialect/ddl.go +++ b/clients/mssql/dialect/ddl.go @@ -46,3 +46,11 @@ func (MSSQLDialect) BuildCreateTableQuery(tableID sql.TableIdentifier, _ bool, c // Microsoft SQL Server doesn't support IF NOT EXISTS return fmt.Sprintf("CREATE TABLE %s (%s);", tableID.FullyQualifiedName(), strings.Join(colSQLParts, ",")) } + +func (MSSQLDialect) BuildDropTableQuery(tableID sql.TableIdentifier) string { + return "DROP TABLE IF EXISTS " + tableID.FullyQualifiedName() +} + +func (MSSQLDialect) BuildTruncateTableQuery(tableID sql.TableIdentifier) string { + return "TRUNCATE TABLE " + tableID.FullyQualifiedName() +} diff --git a/clients/mssql/dialect/dialect_test.go b/clients/mssql/dialect/dialect_test.go index 6d4fdb58a..9ce877b0c 100644 --- a/clients/mssql/dialect/dialect_test.go +++ b/clients/mssql/dialect/dialect_test.go @@ -60,6 +60,20 @@ func TestMSSQLDialect_BuildCreateTableQuery(t *testing.T) { ) } +func TestMSSQLDialect_BuildDropTableQuery(t *testing.T) { + assert.Equal(t, + `DROP TABLE IF EXISTS "schema1"."table1"`, + MSSQLDialect{}.BuildDropTableQuery(NewTableIdentifier("schema1", "table1")), + ) +} + +func TestMSSQLDialect_BuildTruncateTableQuery(t *testing.T) { + assert.Equal(t, + `TRUNCATE TABLE "schema1"."table1"`, + MSSQLDialect{}.BuildTruncateTableQuery(NewTableIdentifier("schema1", "table1")), + ) +} + func TestMSSQLDialect_BuildAddColumnQuery(t *testing.T) { fakeTableID := &mocks.FakeTableIdentifier{} fakeTableID.FullyQualifiedNameReturns("{TABLE}") diff --git a/clients/redshift/dialect/ddl.go b/clients/redshift/dialect/ddl.go index 88e8ccbcd..ad5d8708f 100644 --- a/clients/redshift/dialect/ddl.go +++ b/clients/redshift/dialect/ddl.go @@ -17,25 +17,25 @@ func (RedshiftDialect) BuildDescribeTableQuery(tableID sql.TableIdentifier) (str // This query is a modified fork from: https://gist.github.com/alexanderlz/7302623 return fmt.Sprintf(` -SELECT +SELECT c.column_name, - CASE - WHEN c.data_type = 'numeric' THEN + CASE + WHEN c.data_type = 'numeric' THEN 'numeric(' || COALESCE(CAST(c.numeric_precision AS VARCHAR), '') || ',' || COALESCE(CAST(c.numeric_scale AS VARCHAR), '') || ')' - ELSE + ELSE c.data_type END AS data_type, c.%s, d.description FROM INFORMATION_SCHEMA.COLUMNS c -LEFT JOIN - PG_CLASS c1 ON c.table_name = c1.relname -LEFT JOIN - PG_CATALOG.PG_NAMESPACE n ON c.table_schema = n.nspname AND c1.relnamespace = n.oid -LEFT JOIN - PG_CATALOG.PG_DESCRIPTION d ON d.objsubid = c.ordinal_position AND d.objoid = c1.oid -WHERE +LEFT JOIN + PG_CLASS c1 ON c.table_name = c1.relname +LEFT JOIN + PG_CATALOG.PG_NAMESPACE n ON c.table_schema = n.nspname AND c1.relnamespace = n.oid +LEFT JOIN + PG_CATALOG.PG_DESCRIPTION d ON d.objsubid = c.ordinal_position AND d.objoid = c1.oid +WHERE LOWER(c.table_schema) = LOWER($1) AND LOWER(c.table_name) = LOWER($2); `, constants.StrPrecisionCol), []any{redshiftTableID.Schema(), redshiftTableID.Table()}, nil } @@ -52,3 +52,11 @@ func (RedshiftDialect) BuildCreateTableQuery(tableID sql.TableIdentifier, _ bool // Redshift uses the same syntax for temporary and permanent tables. return fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (%s);", tableID.FullyQualifiedName(), strings.Join(colSQLParts, ",")) } + +func (RedshiftDialect) BuildDropTableQuery(tableID sql.TableIdentifier) string { + return "DROP TABLE IF EXISTS " + tableID.FullyQualifiedName() +} + +func (RedshiftDialect) BuildTruncateTableQuery(tableID sql.TableIdentifier) string { + return "TRUNCATE TABLE " + tableID.FullyQualifiedName() +} diff --git a/clients/redshift/dialect/dialect_test.go b/clients/redshift/dialect/dialect_test.go index 1e7bb8992..1ab2e952f 100644 --- a/clients/redshift/dialect/dialect_test.go +++ b/clients/redshift/dialect/dialect_test.go @@ -45,6 +45,20 @@ func TestRedshiftDialect_BuildCreateTableQuery(t *testing.T) { ) } +func TestRedshiftDialect_BuildDropTableQuery(t *testing.T) { + assert.Equal(t, + `DROP TABLE IF EXISTS schema1."table1"`, + RedshiftDialect{}.BuildDropTableQuery(NewTableIdentifier("schema1", "table1")), + ) +} + +func TestRedshiftDialect_BuildTruncateTableQuery(t *testing.T) { + assert.Equal(t, + `TRUNCATE TABLE schema1."table1"`, + RedshiftDialect{}.BuildTruncateTableQuery(NewTableIdentifier("schema1", "table1")), + ) +} + func TestRedshiftDialect_BuildAddColumnQuery(t *testing.T) { fakeTableID := &mocks.FakeTableIdentifier{} fakeTableID.FullyQualifiedNameReturns("{TABLE}") diff --git a/clients/snowflake/dialect/ddl.go b/clients/snowflake/dialect/ddl.go index ff0b5b699..91c93bf06 100644 --- a/clients/snowflake/dialect/ddl.go +++ b/clients/snowflake/dialect/ddl.go @@ -21,6 +21,14 @@ func (SnowflakeDialect) BuildCreateTableQuery(tableID sql.TableIdentifier, tempo } } +func (SnowflakeDialect) BuildDropTableQuery(tableID sql.TableIdentifier) string { + return "DROP TABLE IF EXISTS " + tableID.FullyQualifiedName() +} + +func (SnowflakeDialect) BuildTruncateTableQuery(tableID sql.TableIdentifier) string { + return "TRUNCATE TABLE IF EXISTS " + tableID.FullyQualifiedName() +} + func (SnowflakeDialect) BuildAddColumnQuery(tableID sql.TableIdentifier, sqlPart string) string { return fmt.Sprintf("ALTER TABLE %s ADD COLUMN IF NOT EXISTS %s", tableID.FullyQualifiedName(), sqlPart) } diff --git a/clients/snowflake/dialect/dialect_test.go b/clients/snowflake/dialect/dialect_test.go index 6f808433a..4fb59e240 100644 --- a/clients/snowflake/dialect/dialect_test.go +++ b/clients/snowflake/dialect/dialect_test.go @@ -46,6 +46,20 @@ func TestSnowflakeDialect_BuildCreateTableQuery(t *testing.T) { ) } +func TestSnowflakeDialect_BuildDropTableQuery(t *testing.T) { + assert.Equal(t, + `DROP TABLE IF EXISTS database1.schema1."TABLE1"`, + SnowflakeDialect{}.BuildDropTableQuery(NewTableIdentifier("database1", "schema1", "table1")), + ) +} + +func TestSnowflakeDialect_BuildTruncateTableQuery(t *testing.T) { + assert.Equal(t, + `TRUNCATE TABLE IF EXISTS database1.schema1."TABLE1"`, + SnowflakeDialect{}.BuildTruncateTableQuery(NewTableIdentifier("database1", "schema1", "table1")), + ) +} + func TestSnowflakeDialect_BuildAddColumnQuery(t *testing.T) { fakeTableID := &mocks.FakeTableIdentifier{} fakeTableID.FullyQualifiedNameReturns("{TABLE}") diff --git a/lib/sql/dialect.go b/lib/sql/dialect.go index bbe624561..f39f5e7d0 100644 --- a/lib/sql/dialect.go +++ b/lib/sql/dialect.go @@ -31,6 +31,8 @@ type Dialect interface { IsColumnAlreadyExistsErr(err error) bool IsTableDoesNotExistErr(err error) bool BuildCreateTableQuery(tableID TableIdentifier, temporary bool, colSQLParts []string) string + BuildDropTableQuery(tableID TableIdentifier) string + BuildTruncateTableQuery(tableID TableIdentifier) string BuildDedupeQueries(tableID, stagingTableID TableIdentifier, primaryKeys []string, includeArtieUpdatedAt bool) []string BuildDedupeTableQuery(tableID TableIdentifier, primaryKeys []string) string BuildDescribeTableQuery(tableID TableIdentifier) (string, []any, error)