kaxil commented on a change in pull request #9973:
URL: https://github.com/apache/airflow/pull/9973#discussion_r629693129



##########
File path: 
airflow/migrations/versions/bbf4a7ad0465_remove_id_column_from_xcom.py
##########
@@ -35,6 +37,61 @@
 depends_on = None
 
 
+def get_table_constraints(conn, table_name):
+    """
+    This function return primary and unique constraint
+    along with column name. some tables like task_instance
+    is missing primary key constraint name and the name is
+    auto-generated by sql server. so this function helps to
+    retrieve any primary or unique constraint name.
+    :param conn: sql connection object
+    :param table_name: table name
+    :return: a dictionary of ((constraint name, constraint type), column name) 
of table
+    :rtype: defaultdict(list)
+    """
+    query = """SELECT tc.CONSTRAINT_NAME , tc.CONSTRAINT_TYPE, ccu.COLUMN_NAME
+     FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc
+     JOIN INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE AS ccu ON 
ccu.CONSTRAINT_NAME = tc.CONSTRAINT_NAME
+     WHERE tc.TABLE_NAME = '{table_name}' AND
+     (tc.CONSTRAINT_TYPE = 'PRIMARY KEY' or UPPER(tc.CONSTRAINT_TYPE) = 
'UNIQUE')
+    """.format(
+        table_name=table_name
+    )
+    result = conn.execute(query).fetchall()
+    constraint_dict = defaultdict(list)
+    for constraint, constraint_type, column in result:
+        constraint_dict[(constraint, constraint_type)].append(column)
+    return constraint_dict
+
+
+def drop_column_constraints(operator, column_name, constraint_dict):
+    """
+    Drop a primary key or unique constraint
+    :param operator: batch_alter_table for the table

Review comment:
       ```suggestion
       Drop a primary key or unique constraint
   
       :param operator: batch_alter_table for the table
   ```




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


Reply via email to