Getting Snowflake Primary Key Columns as a Table
One of my customers had an interesting requirement. In order to dynamically create merge statements, they needed a way to collect the primary key columns for any given table. After discussing some options — returning them as a delimited string, array, etc., we agreed that returning the columns in a table is the best option.
This User Defined Table Function (UDTF) returns the columns for a table’s primary key. The UDTF will return a table with a single column, each row in the table is one of the columns in the input table’s primary key. If there is no primary key, the table will have no rows. For a single-column primary key, the table will have the one row, and for composite primary keys it will return all columns in the key.
One thing you may notice is that the input to the UDTF is the table’s DDL, not just the table’s name. The reason for this is because UDTFs cannot execute SQL. The simplest way to handle this situation is to nest the GET_DDL function as the parameter for the GET_PK_COLUMNS function. You can see how this works in the code samp
create database DB_Primary_Keys;
/**************************************************************************************************************
* *
* Set up test tables with four types of primary key: Named composite, unnamed composite, inline, and none. *
* *
**************************************************************************************************************/
-- Named multi-column PK
create or replace temporary table table1
(
column_name1 number NOT NULL,
column_name2 number NOT NULL,
column_name3 string,
CONSTRAINT Constraint_name PRIMARY KEY (column_name1, column_name2)
);
-- Unnamed multi-column PK
create temporary table table2
(
column_name1 number NOT NULL,
column_name2 number NOT NULL,
column_name3 string,
PRIMARY KEY (column_name1, column_name2)
);
-- Inline single-column PK
create or replace temporary table table3
(
column_name1 number primary key,
column_name2 number NOT NULL,
column_name3 string
);
-- No PK defined
create or replace temporary table table4
(
column_name1 number,
column_name2 number,
column_name3 string
);
/********************************************************************************************************
* *
* User defined table function (UDTF) to get primary keys for a table. *
* *
* @param {string}: TABLE_DDL The DDL for the table to get the PKs. Usually use get_ddl. *
* @return {table}: A table with the columns comprising the table's primary key *
* *
********************************************************************************************************/
create or replace function GET_PK_COLUMNS(TABLE_DDL string)
returns table (PK_COLUMN string)
language javascript
as
$$
{
processRow: function get_params(row, rowWriter, context){
var pkCols = getPKs(row.TABLE_DDL);
for (i = 0; i < pkCols.length; i++) {
rowWriter.writeRow({PK_COLUMN: pkCols[i]});
}
function getPKs(tableDDL) {
var c;
var keyword = "primary key";
var ins = -1;
var s = tableDDL.split("\n");
for (var i = 0; i < s.length; i++) {
ins = s[i].indexOf(keyword);
if (ins != -1) {
var colList = s[i].substring(ins + keyword.length);
colList = colList.replace("(", "");
colList = colList.replace(")", "");
var colArray = colList.split(",");
for (pkc = 0; c < colArray.length; pkc++) {
colArray[pkc] = colArray[pkc].trim();
}
return colArray;
}
}
return []; // No PK
}
}
}
$$;
/**************************************************************************************************************
* *
* Test execution of the UDTF. *
* *
**************************************************************************************************************/
select * from table(get_pk_columns(get_ddl('table', 'table1'))) PKS; -- Multi-column PK with named constraint
select * from table(get_pk_columns(get_ddl('table', 'table2'))) PKS; -- Multi-column PK with no name for constraint
select * from table(get_pk_columns(get_ddl('table', 'table3'))) PKS; -- Single column PK inline definition
select * from table(get_pk_columns(get_ddl('table', 'table4'))) PKS; -- No PKs